classify.py

#!/usr/bin/env python3
"""
Metal Detector Audio Classifier

Classify individual audio files using trained advanced models.

Usage:
    python classify.py audio_file.wav
    python classify.py --help
"""

import argparse
import sys
from pathlib import Path
import json

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

from src.ml.advanced_classifier import AdvancedMetalClassifier

console = Console()

def display_classification_result(result, audio_file):
    """Display classification results in a rich format."""
    console.print(Panel.fit(f"🔍 Metal Detector Classification: {audio_file.name}", style="bold blue"))
    
    # Main prediction
    confidence_color = "green" if result.confidence > 0.8 else "yellow" if result.confidence > 0.6 else "red"
    console.print(f"🎯 **Predicted Label**: [{confidence_color}]{result.predicted_label.upper()}[/{confidence_color}]")
    console.print(f"🔥 **Confidence**: [{confidence_color}]{result.confidence:.3f}[/{confidence_color}]")
    console.print(f"⚡ **Processing Time**: {result.processing_time:.3f} seconds")
    
    # All probabilities
    console.print("\n📊 **All Label Probabilities**:")
    prob_table = Table(show_header=True)
    prob_table.add_column("Label", style="cyan")
    prob_table.add_column("Probability", style="magenta", justify="right")
    prob_table.add_column("Percentage", style="yellow", justify="right")
    
    for label, prob in sorted(result.all_probabilities.items(), key=lambda x: x[1], reverse=True):
        prob_table.add_row(
            label.title(),
            f"{prob:.4f}",
            f"{prob*100:.1f}%"
        )
    
    console.print(prob_table)
    
    # Model ensemble votes
    console.print("\n🤖 **Model Ensemble Votes**:")
    ensemble_table = Table(show_header=True)
    ensemble_table.add_column("Model", style="cyan")
    ensemble_table.add_column("Prediction", style="magenta")
    ensemble_table.add_column("Agreement", style="green")
    
    main_prediction = result.predicted_label
    for model, prediction in result.model_ensemble_votes.items():
        agreement = "✅ Agrees" if prediction == main_prediction else "❌ Disagrees"
        ensemble_table.add_row(
            model.upper(),
            prediction.title(),
            agreement
        )
    
    console.print(ensemble_table)
    
    # Feature importance (top 10)
    if result.feature_importance:
        console.print("\n🔬 **Top Features Contributing to Classification**:")
        feature_table = Table(show_header=True)
        feature_table.add_column("Feature", style="cyan")
        feature_table.add_column("Importance", style="magenta", justify="right")
        
        top_features = list(result.feature_importance.items())[:10]
        for feature, importance in top_features:
            feature_table.add_row(feature.replace('_', ' ').title(), f"{importance:.4f}")
        
        console.print(feature_table)
    
    # Spectral analysis summary
    if result.spectral_analysis:
        console.print("\n🌊 **Spectral Analysis Summary**:")
        spectral_info = []
        for key, value in result.spectral_analysis.items():
            if isinstance(value, (int, float)):
                spectral_info.append(f"**{key.replace('_', ' ').title()}**: {value:.2f}")
        
        if spectral_info:
            console.print(" | ".join(spectral_info[:5]))  # Show first 5
    
    # Temporal analysis summary
    if result.temporal_analysis:
        console.print("\n⏱️ **Temporal Analysis Summary**:")
        temporal_info = []
        for key, value in result.temporal_analysis.items():
            if isinstance(value, (int, float)):
                temporal_info.append(f"**{key.replace('_', ' ').title()}**: {value:.2f}")
        
        if temporal_info:
            console.print(" | ".join(temporal_info[:5]))  # Show first 5

def main():
    """Main classification function."""
    parser = argparse.ArgumentParser(description="Classify metal detector audio files")
    parser.add_argument("audio_file", type=str, help="Path to audio file to classify (supports .wav, .mp3, .m4a)")
    parser.add_argument("--model-dir", type=str, default="models/advanced",
                       help="Directory containing trained models")
    parser.add_argument("--output", type=str, help="Save results to JSON file")
    parser.add_argument("--verbose", action="store_true", help="Show detailed analysis")
    
    args = parser.parse_args()
    
    audio_file = Path(args.audio_file)
    model_dir = Path(args.model_dir)
    
    # Validate inputs
    if not audio_file.exists():
        console.print(f"❌ Audio file not found: {audio_file}")
        sys.exit(1)
    
    if not model_dir.exists():
        console.print(f"❌ Model directory not found: {model_dir}")
        console.print("💡 Train models first using: python train_model.py")
        sys.exit(1)
    
    # Initialize classifier
    console.print("🤖 Loading advanced metal detector classifier...")
    
    try:
        with Progress() as progress:
            task = progress.add_task("Loading models...", total=100)
            
            classifier = AdvancedMetalClassifier(model_path=model_dir)
            progress.update(task, advance=30)
            
            # Check if models exist
            if not classifier._model_exists():
                console.print("❌ Trained models not found!")
                console.print("💡 Train models first using: python train_model.py")
                sys.exit(1)
            
            # Load models
            classifier.load_model()
            progress.update(task, advance=70)
            
            console.print("✅ Models loaded successfully!")
        
        # Classify the audio file
        console.print(f"🎵 Analyzing: {audio_file.name}")
        
        with console.status("[bold green]Classifying audio...") as status:
            result = classifier.classify_audio_file(audio_file)
        
        # Display results
        display_classification_result(result, audio_file)
        
        # Save to JSON if requested
        if args.output:
            output_path = Path(args.output)
            
            # Convert result to JSON-serializable format
            result_dict = {
                'audio_file': str(audio_file),
                'predicted_label': result.predicted_label,
                'confidence': result.confidence,
                'all_probabilities': result.all_probabilities,
                'processing_time': result.processing_time,
                'model_ensemble_votes': result.model_ensemble_votes,
                'feature_importance': result.feature_importance,
                'spectral_analysis': result.spectral_analysis,
                'temporal_analysis': result.temporal_analysis
            }
            
            with open(output_path, 'w') as f:
                json.dump(result_dict, f, indent=2, default=str)
            
            console.print(f"💾 Results saved to: {output_path}")
        
        # Classification summary
        confidence_level = "HIGH" if result.confidence > 0.8 else "MEDIUM" if result.confidence > 0.6 else "LOW"
        console.print(f"\n🎯 **Final Result**: {result.predicted_label.upper()} ({confidence_level} confidence)")
        
    except Exception as e:
        console.print(f"❌ Classification failed: {e}")
        if args.verbose:
            import traceback
            console.print(traceback.format_exc())
        sys.exit(1)

if __name__ == "__main__":
    main()