Remove redundant cursor styling

clio-ai/real_time/train_classifier.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""
+Train an audio embedding classifier on recorded samples.
+
+Uses wav2vec2 to extract embeddings, then trains a simple classifier.
+
+Usage:
+    uv run python train_classifier.py
+"""
+
+import pickle
+from pathlib import Path
+
+import numpy as np
+import torch
+from scipy.io import wavfile
+from torchaudio.pipelines import WAV2VEC2_BASE
+
+RECORDINGS_DIR = Path(__file__).parent / "recordings"
+MODEL_PATH = Path(__file__).parent / "classifier.pkl"
+SAMPLE_RATE = 16000
+
+# Categories (excluding wake_word which is handled separately)
+CATEGORIES = [
+    "pressure_below_15",
+    "pressure_15_to_25",
+    "pressure_25_to_35",
+    "pressure_above_35",
+    "drop_l",
+    "drop_p",
+    "drop_d",
+    "both_eyes",
+]
+
+
+def load_audio(path: Path) -> torch.Tensor:
+    """Load and preprocess audio file."""
+    sr, data = wavfile.read(path)
+
+    # Convert to float32
+    if data.dtype == np.int16:
+        data = data.astype(np.float32) / 32768.0
+    elif data.dtype == np.int32:
+        data = data.astype(np.float32) / 2147483648.0
+    elif data.dtype != np.float32:
+        data = data.astype(np.float32)
+
+    # Convert to mono if stereo
+    if len(data.shape) > 1:
+        data = data.mean(axis=1)
+
+    # Resample if needed
+    if sr != SAMPLE_RATE:
+        from scipy import signal
+        num_samples = int(len(data) * SAMPLE_RATE / sr)
+        data = signal.resample(data, num_samples)
+
+    # Convert to torch tensor with batch dimension
+    waveform = torch.from_numpy(data).unsqueeze(0)
+    return waveform
+
+
+def extract_embedding(waveform: torch.Tensor, model, device: str) -> np.ndarray:
+    """Extract embedding from audio using wav2vec2."""
+    with torch.no_grad():
+        waveform = waveform.to(device)
+        features, _ = model.extract_features(waveform)
+        # Use the last layer's features, averaged over time
+        embedding = features[-1].mean(dim=1).cpu().numpy()
+    return embedding.flatten()
+
+
+def load_samples():
+    """Load all training samples."""
+    samples = []
+    labels = []
+
+    for category in CATEGORIES:
+        category_dir = RECORDINGS_DIR / category
+        if not category_dir.exists():
+            print(f"Warning: {category_dir} does not exist")
+            continue
+
+        wav_files = list(category_dir.glob("*.wav"))
+        print(f"  {category}: {len(wav_files)} samples")
+
+        for wav_file in wav_files:
+            samples.append(wav_file)
+            labels.append(category)
+
+    return samples, labels
+
+
+def main():
+    print("=== Training Audio Classifier ===\n")
+
+    # Check for samples
+    if not RECORDINGS_DIR.exists():
+        print(f"Error: {RECORDINGS_DIR} does not exist")
+        print("Run ./collect_samples.sh first to record training samples")
+        return
+
+    print("Loading samples...")
+    sample_paths, labels = load_samples()
+
+    if len(sample_paths) == 0:
+        print("\nNo samples found. Run ./collect_samples.sh first.")
+        return
+
+    if len(set(labels)) < 2:
+        print(f"\nNeed samples from at least 2 categories. Found: {set(labels)}")
+        return
+
+    print(f"\nTotal: {len(sample_paths)} samples across {len(set(labels))} categories")
+
+    # Check minimum samples per category
+    from collections import Counter
+    counts = Counter(labels)
+    min_samples = min(counts.values())
+    if min_samples < 3:
+        print(f"\nWarning: Some categories have < 3 samples. More samples = better accuracy.")
+
+    # Load model
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"\nLoading wav2vec2 on {device}...")
+    bundle = WAV2VEC2_BASE
+    model = bundle.get_model().to(device)
+    model.eval()
+
+    # Extract embeddings
+    print("Extracting embeddings...")
+    embeddings = []
+    valid_labels = []
+
+    for i, (path, label) in enumerate(zip(sample_paths, labels)):
+        try:
+            waveform = load_audio(path)
+            embedding = extract_embedding(waveform, model, device)
+            embeddings.append(embedding)
+            valid_labels.append(label)
+            print(f"  [{i+1}/{len(sample_paths)}] {path.name}")
+        except Exception as e:
+            print(f"  [{i+1}/{len(sample_paths)}] {path.name} - ERROR: {e}")
+
+    if len(embeddings) < 2:
+        print("\nNot enough valid samples to train.")
+        return
+
+    X = np.array(embeddings)
+    y = np.array(valid_labels)
+
+    print(f"\nEmbedding shape: {X.shape}")
+
+    # Train classifier
+    print("\nTraining classifier...")
+    from sklearn.preprocessing import LabelEncoder
+    from sklearn.svm import SVC
+    from sklearn.model_selection import cross_val_score
+
+    label_encoder = LabelEncoder()
+    y_encoded = label_encoder.fit_transform(y)
+
+    # Use SVM with probability estimates
+    classifier = SVC(kernel='rbf', probability=True, C=1.0)
+
+    # Cross-validation if enough samples
+    if len(X) >= 4:
+        n_splits = min(5, len(X))
+        scores = cross_val_score(classifier, X, y_encoded, cv=n_splits)
+        print(f"Cross-validation accuracy: {scores.mean():.2%} (+/- {scores.std()*2:.2%})")
+
+    # Train on all data
+    classifier.fit(X, y_encoded)
+
+    # Save model
+    model_data = {
+        'classifier': classifier,
+        'label_encoder': label_encoder,
+        'categories': CATEGORIES,
+    }
+
+    with open(MODEL_PATH, 'wb') as f:
+        pickle.dump(model_data, f)
+
+    print(f"\nClassifier saved to: {MODEL_PATH}")
+    print(f"Categories: {list(label_encoder.classes_)}")
+    print("\nDone! The listener will now use this classifier.")
+
+
+if __name__ == "__main__":
+    main()