stream_classify.py

#!/usr/bin/env python3
"""
Real-time Metal Detector Audio Classification

Stream live audio from metal detectors and classify patterns in real-time
using advanced machine learning models.

Usage:
    python stream_classify.py --device 0 --duration 60
    python stream_classify.py --list-devices
    python stream_classify.py --help
"""

import argparse
import sys
from pathlib import Path
import time

from rich.console import Console
from rich.table import Table
from rich.panel import Panel

from src.audio.streaming import RealTimeMetalDetector, StreamingConfig, list_audio_devices

console = Console()

def main():
    """Main real-time classification function."""
    parser = argparse.ArgumentParser(description="Real-time metal detector audio classification")
    parser.add_argument("--device", type=int, help="Audio input device ID")
    parser.add_argument("--list-devices", action="store_true", help="List available audio devices")
    parser.add_argument("--duration", type=float, default=60.0, help="Recording duration in seconds")
    parser.add_argument("--confidence-threshold", type=float, default=0.7, help="Confidence threshold for detection")
    parser.add_argument("--classification-interval", type=float, default=2.0, help="Seconds between classifications")
    parser.add_argument("--save-detections", type=str, help="Directory to save detected patterns")
    parser.add_argument("--model-dir", type=str, default="models/advanced", help="Model directory")
    parser.add_argument("--no-visualization", action="store_true", help="Disable real-time visualization")
    
    args = parser.parse_args()
    
    # List devices if requested
    if args.list_devices:
        devices = list_audio_devices()
        console.print("🎤 Available Audio Input Devices:", style="bold")
        
        if not devices:
            console.print("❌ No audio input devices found!")
            return
        
        table = Table()
        table.add_column("ID", style="cyan")
        table.add_column("Device Name", style="green")
        table.add_column("Channels", style="yellow")
        table.add_column("Sample Rate", style="blue")
        
        for device in devices:
            table.add_row(
                str(device['id']),
                device['name'],
                str(device['channels']),
                f"{device['sample_rate']:.0f} Hz"
            )
        
        console.print(table)
        return
    
    # Initialize streaming configuration
    config = StreamingConfig(
        device_id=args.device,
        confidence_threshold=args.confidence_threshold,
        classification_interval=args.classification_interval
    )
    
    model_dir = Path(args.model_dir)
    
    console.print(Panel.fit("🔍 Real-time Metal Detector AI", style="bold blue"))
    
    # Initialize classifier
    try:
        detector = RealTimeMetalDetector(config=config, model_path=model_dir)
    except Exception as e:
        console.print(f"❌ Error initializing detector: {e}", style="red")
        return
    
    # Check if models exist
    if not detector.classifier.is_ready():
        console.print("❌ No trained models found!")
        console.print("💡 Train models first using: python train_model.py")
        return
    
    # Display configuration
    console.print(f"🎤 Audio Device: {args.device or 'default'}")
    console.print(f"⏱️ Duration: {args.duration} seconds") 
    console.print(f"🎯 Classification Interval: {args.classification_interval}s")
    console.print(f"🔥 Confidence Threshold: {args.confidence_threshold}")
    console.print(f"📁 Model Directory: {model_dir}")
    
    # Start streaming
    try:
        detector.start_streaming(visualize=not args.no_visualization)
        
        if args.no_visualization:
            console.print(f"🎤 Streaming started... Running for {args.duration} seconds")
            console.print("Press Ctrl+C to stop early")
        
        # Run for specified duration
        time.sleep(args.duration)
        
    except KeyboardInterrupt:
        console.print("\n⏹️ Stopped by user", style="yellow")
    except Exception as e:
        console.print(f"❌ Error during streaming: {e}", style="red")
    finally:
        detector.stop_streaming()
        
        # Show final statistics
        stats = detector.get_detection_summary()
        console.print("\n📊 Final Session Statistics:", style="bold")
        console.print(f"Total detections: {stats['total_detections']}")
        console.print(f"High confidence: {stats.get('high_confidence_detections', 0)}")
        console.print(f"Average confidence: {stats.get('average_confidence', 0):.3f}")
        console.print(f"Session duration: {stats.get('session_duration', 0):.1f}s")
        
        # Show label distribution
        if 'label_distribution' in stats and stats['label_distribution']:
            console.print("\n🏷️ Label Distribution:")
            label_table = Table()
            label_table.add_column("Label", style="cyan")
            label_table.add_column("Count", style="magenta")
            label_table.add_column("Percentage", style="yellow")
            
            total_labels = sum(stats['label_distribution'].values())
            for label, count in sorted(stats['label_distribution'].items(), key=lambda x: x[1], reverse=True):
                percentage = (count / total_labels) * 100 if total_labels > 0 else 0
                label_table.add_row(
                    label.title(),
                    str(count),
                    f"{percentage:.1f}%"
                )
            
            console.print(label_table)
        
        # Save detections if requested
        if args.save_detections:
            output_path = Path(args.save_detections)
            output_path.mkdir(parents=True, exist_ok=True)
            
            session_file = output_path / f"session_{int(time.time())}.json"
            detector.save_session(session_file)
            
            console.print(f"💾 Session data saved to: {session_file}")

if __name__ == "__main__":
    main()