classify_audio.py

#!/usr/bin/env python3
"""
Audio Genre Classifier using Fine-Tuned AST Model

This script takes an audio file as input and uses the fine-tuned Audio Spectrogram
Transformer model to classify its genre. It returns the top predictions with
confidence scores.
"""

import argparse
import sys
import numpy as np
import librosa
import torch
from pathlib import Path
from transformers import AutoProcessor, AutoModelForAudioClassification


# --- Configuration ---
MODEL_PATH = "./ast-finetuned-all-styles"
TARGET_SAMPLE_RATE = 16000
MAX_DURATION_S = 10.0  # Match training duration


def load_model_and_processor(model_path):
    """Load the fine-tuned model and processor."""
    print(f"Loading model from {model_path}...", file=sys.stderr)
    
    processor = AutoProcessor.from_pretrained(model_path)
    model = AutoModelForAudioClassification.from_pretrained(model_path)
    
    # Set to evaluation mode
    model.eval()
    
    print(f"Model loaded successfully. Classes: {len(model.config.id2label)}", file=sys.stderr)
    return model, processor


def preprocess_audio(file_path, processor):
    """Load and preprocess an audio file for classification."""
    print(f"Loading audio from {file_path}...", file=sys.stderr)
    
    # Load audio file
    y, sr = librosa.load(file_path, sr=TARGET_SAMPLE_RATE, mono=True)
    
    # Trim silence
    yt, _ = librosa.effects.trim(y, top_db=60, ref=1.0)
    y = yt
    
    # Truncate or pad to MAX_DURATION_S
    total_samples = int(MAX_DURATION_S * TARGET_SAMPLE_RATE)
    if len(y) > total_samples:
        # Take a middle section for consistency
        start_idx = (len(y) - total_samples) // 2
        y = y[start_idx:start_idx + total_samples]
    else:
        # Pad with zeros
        y = np.pad(y, (0, total_samples - len(y)), 'constant')
    
    # Process with the processor (creates mel spectrogram)
    inputs = processor(
        y,
        sampling_rate=TARGET_SAMPLE_RATE,
        return_tensors="pt"
    )
    
    # Permute dimensions to match training format
    # (batch_size, num_frames, num_mel_bins) -> (batch_size, num_mel_bins, num_frames)
    inputs["input_values"] = inputs["input_values"].permute(0, 2, 1)
    
    return inputs


def classify_audio(file_path, model, processor, top_k=5):
    """Classify an audio file and return top predictions."""
    
    # Preprocess audio
    inputs = preprocess_audio(file_path, processor)
    
    # Run inference
    print("Running classification...", file=sys.stderr)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    # Get probabilities using softmax
    probs = torch.nn.functional.softmax(logits, dim=-1)
    
    # Get top k predictions
    top_probs, top_indices = torch.topk(probs[0], k=min(top_k, len(probs[0])))
    
    # Build results
    results = []
    for prob, idx in zip(top_probs, top_indices):
        genre = model.config.id2label[idx.item()]
        confidence = prob.item()
        results.append({
            'genre': genre,
            'confidence': confidence,
            'label_id': idx.item()
        })
    
    return results


def main():
    parser = argparse.ArgumentParser(
        description="Classify audio file genre using fine-tuned AST model"
    )
    parser.add_argument(
        "audio_file",
        type=str,
        help="Path to the audio file to classify"
    )
    parser.add_argument(
        "--model-path",
        type=str,
        default=MODEL_PATH,
        help=f"Path to the fine-tuned model (default: {MODEL_PATH})"
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=5,
        help="Number of top predictions to return (default: 5)"
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output results in JSON format"
    )
    
    args = parser.parse_args()
    
    # Check if file exists
    audio_path = Path(args.audio_file)
    if not audio_path.exists():
        print(f"Error: File not found: {args.audio_file}", file=sys.stderr)
        sys.exit(1)
    
    # Load model
    try:
        model, processor = load_model_and_processor(args.model_path)
    except Exception as e:
        print(f"Error loading model: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Classify audio
    try:
        results = classify_audio(args.audio_file, model, processor, args.top_k)
    except Exception as e:
        print(f"Error classifying audio: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Output results
    if args.json:
        import json
        output = {
            'file': str(audio_path),
            'predictions': results
        }
        print(json.dumps(output, indent=2))
    else:
        print(f"\n{'='*70}")
        print(f"Genre Classification Results for: {audio_path.name}")
        print(f"{'='*70}\n")
        
        for i, result in enumerate(results, 1):
            confidence_pct = result['confidence'] * 100
            bar_length = int(confidence_pct / 2)  # Scale to 50 chars max
            bar = '█' * bar_length + '░' * (50 - bar_length)
            
            print(f"{i}. {result['genre']:<30} {confidence_pct:>6.2f}% {bar}")
        
        print(f"\n{'='*70}")
        print(f"Top prediction: {results[0]['genre']} ({results[0]['confidence']*100:.2f}%)")
        print(f"{'='*70}\n")


if __name__ == "__main__":
    main()