train_andy.py

#!/usr/bin/env python3
"""
Train custom "Andy" wake word model for OpenWakeWord

Two methods:
1. Synthetic: Generate TTS samples of "Andy" (fast but less accurate)
2. Real: Record yourself saying "Andy" multiple times (more accurate)

This script helps with both approaches.
"""

import os
import sys
import json
import wave
import struct
import numpy as np
from pathlib import Path

SAMPLE_RATE = 16000
WAKE_WORD = "andy"
OUTPUT_DIR = Path("~/projects/voice-control/training_data").expanduser()
MODEL_OUTPUT = Path("~/projects/voice-control/models").expanduser()

def setup_directories():
    """Create necessary directories"""
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    MODEL_OUTPUT.mkdir(parents=True, exist_ok=True)
    (OUTPUT_DIR / "positive").mkdir(exist_ok=True)
    (OUTPUT_DIR / "negative").mkdir(exist_ok=True)
    print(f"Created directories in {OUTPUT_DIR}")

def record_sample(filename: str, duration: float = 2.0):
    """Record a single audio sample"""
    import pyaudio
    
    p = pyaudio.PyAudio()
    stream = p.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=SAMPLE_RATE,
        input=True,
        frames_per_buffer=1024
    )
    
    print(f"Recording for {duration} seconds... Say 'Andy'!")
    frames = []
    for _ in range(0, int(SAMPLE_RATE / 1024 * duration)):
        data = stream.read(1024)
        frames.append(data)
    
    stream.stop_stream()
    stream.close()
    p.terminate()
    
    # Save as WAV
    wf = wave.open(str(filename), 'wb')
    wf.setnchannels(1)
    wf.setsampwidth(2)
    wf.setframerate(SAMPLE_RATE)
    wf.writeframes(b''.join(frames))
    wf.close()
    
    print(f"Saved to {filename}")

def record_positive_samples(count: int = 20):
    """Record multiple samples of saying 'Andy'"""
    positive_dir = OUTPUT_DIR / "positive"
    
    print(f"\n=== Recording {count} positive samples ===")
    print("Say 'Andy' clearly after each prompt.\n")
    
    for i in range(count):
        filename = positive_dir / f"andy_{i:03d}.wav"
        input(f"Press Enter to record sample {i+1}/{count}...")
        record_sample(str(filename), duration=1.5)
    
    print(f"\nRecorded {count} positive samples!")

def record_negative_samples(count: int = 20):
    """Record negative samples (random speech, silence, noise)"""
    negative_dir = OUTPUT_DIR / "negative"
    
    print(f"\n=== Recording {count} negative samples ===")
    print("Say random words, make noise, or stay silent.\n")
    
    for i in range(count):
        filename = negative_dir / f"negative_{i:03d}.wav"
        input(f"Press Enter to record negative sample {i+1}/{count}...")
        record_sample(str(filename), duration=2.0)
    
    print(f"\nRecorded {count} negative samples!")

def generate_synthetic_samples():
    """Generate synthetic 'Andy' samples using TTS"""
    try:
        from TTS.api import TTS
    except ImportError:
        print("Installing TTS library...")
        os.system("pip install TTS")
        from TTS.api import TTS
    
    positive_dir = OUTPUT_DIR / "positive"
    
    print("Generating synthetic 'Andy' samples...")
    tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
    
    # Generate variations
    phrases = [
        "Andy", "andy", "Andy!", "Hey Andy", "Okay Andy",
        "Andy?", "ANDY", "andi", "Andie"
    ]
    
    for i, phrase in enumerate(phrases):
        filename = positive_dir / f"synthetic_andy_{i:03d}.wav"
        tts.tts_to_file(text=phrase, file_path=str(filename))
        print(f"Generated: {filename}")
    
    print(f"\nGenerated {len(phrases)} synthetic samples!")

def train_model():
    """Train the OpenWakeWord model"""
    print("\n=== Training Custom Wake Word Model ===\n")
    
    # Check if we have enough samples
    positive_count = len(list((OUTPUT_DIR / "positive").glob("*.wav")))
    negative_count = len(list((OUTPUT_DIR / "negative").glob("*.wav")))
    
    print(f"Positive samples: {positive_count}")
    print(f"Negative samples: {negative_count}")
    
    if positive_count < 10:
        print("\n⚠️  Need at least 10 positive samples!")
        print("Run: python train_andy.py record")
        return
    
    # OpenWakeWord training
    print("\nTraining model... (this may take a few minutes)")
    
    # Use openwakeword's training functionality
    from openwakeword.train import train_model as oww_train
    
    model_path = MODEL_OUTPUT / "andy_v1.onnx"
    
    oww_train(
        positive_audio_dir=str(OUTPUT_DIR / "positive"),
        negative_audio_dir=str(OUTPUT_DIR / "negative") if negative_count > 0 else None,
        output_path=str(model_path),
        epochs=100
    )
    
    print(f"\n✅ Model saved to: {model_path}")
    print("\nTo use this model, update server.py:")
    print(f'  WAKE_WORD_MODELS = ["{model_path}"]')

def main():
    setup_directories()
    
    if len(sys.argv) < 2:
        print("""
Usage: python train_andy.py <command>

Commands:
  record    - Record real samples (recommended, ~5 minutes)
  synthetic - Generate synthetic samples using TTS (faster)
  train     - Train the model after recording samples
  
Recommended workflow:
  1. python train_andy.py record
  2. python train_andy.py train
""")
        return
    
    command = sys.argv[1]
    
    if command == "record":
        record_positive_samples(20)
        print("\nDo you want to record negative samples too? (recommended)")
        if input("Record negative samples? [y/N]: ").lower() == 'y':
            record_negative_samples(10)
    
    elif command == "synthetic":
        generate_synthetic_samples()
    
    elif command == "train":
        train_model()
    
    else:
        print(f"Unknown command: {command}")

if __name__ == "__main__":
    main()