segment_dataset.py

#!/usr/bin/env python3
"""
Segment Dataset for Metal Detector AI

Processes existing audio files to detect and segment individual metal detection events.
This creates multiple training samples from continuous recordings that contain
multiple detection events.

Usage:
    python segment_dataset.py --input-dir data --output-dir data_segmented
    python segment_dataset.py --in-place  # Updates existing dataset structure
"""

import argparse
import shutil
from pathlib import Path
import numpy as np
import librosa
import soundfile as sf
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from rich.table import Table
from typing import Dict, List

from src.data.event_detector import MetalDetectionEventSegmenter

console = Console()


def process_audio_file(file_path: Path, segmenter: MetalDetectionEventSegmenter) -> Dict[str, any]:
    """
    Process a single audio file to detect and segment events.
    
    Args:
        file_path: Path to audio file
        segmenter: Event segmenter instance
        
    Returns:
        Dictionary with processing results
    """
    try:
        # Load audio
        audio, sr = librosa.load(file_path, sr=22050)
        duration = len(audio) / sr
        
        # Detect events
        events = segmenter.detect_events(audio)
        
        # Segment audio
        segments = segmenter.segment_audio(audio, events)
        
        return {
            'success': True,
            'duration': duration,
            'events': events,
            'segments': segments,
            'error': None
        }
        
    except Exception as e:
        return {
            'success': False,
            'duration': 0,
            'events': [],
            'segments': [],
            'error': str(e)
        }


def save_segments(segments: List[Dict], output_dir: Path, base_name: str, label: str) -> int:
    """
    Save audio segments to disk.
    
    Args:
        segments: List of segment dictionaries
        output_dir: Output directory
        base_name: Base filename for segments
        label: Metal type label
        
    Returns:
        Number of segments saved
    """
    saved_count = 0
    
    for i, segment in enumerate(segments):
        # Skip background segments if they're too quiet
        if segment['type'] == 'background':
            rms = np.sqrt(np.mean(segment['audio'] ** 2))
            if rms < 0.001:  # Very quiet, skip
                continue
        
        # Create filename
        filename = f"{base_name}_seg{i:03d}_{segment['type']}.wav"
        output_path = output_dir / filename
        
        # Save audio
        sf.write(output_path, segment['audio'], 22050)
        saved_count += 1
        
    return saved_count


def process_dataset(input_dir: Path, output_dir: Path, in_place: bool = False) -> None:
    """
    Process entire dataset to segment detection events.
    
    Args:
        input_dir: Input dataset directory
        output_dir: Output directory for segmented data
        in_place: If True, add segments to existing structure
    """
    console.print(f"\n[bold blue]Metal Detector AI - Dataset Segmentation[/bold blue]")
    console.print(f"Input directory: {input_dir}")
    console.print(f"Output directory: {output_dir}")
    console.print(f"Mode: {'In-place' if in_place else 'Copy to new directory'}\n")
    
    # Initialize segmenter
    segmenter = MetalDetectionEventSegmenter()
    
    # Find all class directories
    class_dirs = [d for d in input_dir.iterdir() if d.is_dir() and not d.name.startswith('.')]
    
    if not class_dirs:
        console.print("[red]No class directories found in input directory![/red]")
        return
    
    # Statistics
    total_files = 0
    total_events = 0
    total_segments = 0
    class_stats = {}
    
    # Process each class
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        console=console
    ) as progress:
        
        for class_dir in class_dirs:
            label = class_dir.name
            
            # Find audio files
            audio_files = list(class_dir.glob('*.wav')) + \
                         list(class_dir.glob('*.mp3')) + \
                         list(class_dir.glob('*.m4a'))
            
            if not audio_files:
                continue
            
            # Create output directory
            if in_place:
                class_output_dir = class_dir
            else:
                class_output_dir = output_dir / label
                class_output_dir.mkdir(parents=True, exist_ok=True)
            
            # Process files
            task = progress.add_task(f"Processing {label}", total=len(audio_files))
            
            class_files = 0
            class_events = 0
            class_segments = 0
            
            for audio_file in audio_files:
                # Skip if it's already a segment
                if '_seg' in audio_file.stem:
                    progress.advance(task)
                    continue
                
                # Process file
                result = process_audio_file(audio_file, segmenter)
                
                if result['success']:
                    # Save segments
                    saved = save_segments(
                        result['segments'],
                        class_output_dir,
                        audio_file.stem,
                        label
                    )
                    
                    class_files += 1
                    class_events += len(result['events'])
                    class_segments += saved
                    
                    # Copy original file if not in-place
                    if not in_place:
                        shutil.copy2(audio_file, class_output_dir / audio_file.name)
                else:
                    console.print(f"[yellow]Warning: Failed to process {audio_file.name}: {result['error']}[/yellow]")
                
                progress.advance(task)
            
            # Update statistics
            total_files += class_files
            total_events += class_events
            total_segments += class_segments
            
            class_stats[label] = {
                'files': class_files,
                'events': class_events,
                'segments': class_segments
            }
    
    # Display summary
    console.print("\n[bold green]Segmentation Complete![/bold green]\n")
    
    # Create summary table
    table = Table(title="Dataset Segmentation Summary")
    table.add_column("Class", style="cyan")
    table.add_column("Files", justify="right")
    table.add_column("Events Detected", justify="right")
    table.add_column("Segments Created", justify="right")
    table.add_column("Avg Events/File", justify="right")
    
    for label, stats in class_stats.items():
        avg_events = stats['events'] / stats['files'] if stats['files'] > 0 else 0
        table.add_row(
            label,
            str(stats['files']),
            str(stats['events']),
            str(stats['segments']),
            f"{avg_events:.1f}"
        )
    
    # Add totals
    table.add_section()
    avg_total = total_events / total_files if total_files > 0 else 0
    table.add_row(
        "TOTAL",
        str(total_files),
        str(total_events),
        str(total_segments),
        f"{avg_total:.1f}",
        style="bold"
    )
    
    console.print(table)
    
    # Additional insights
    console.print("\n[bold]Insights:[/bold]")
    console.print(f"• Original files: {total_files}")
    console.print(f"• Detection events found: {total_events}")
    console.print(f"• Training segments created: {total_segments}")
    console.print(f"• Average segments per file: {total_segments/total_files:.1f}")
    
    if total_events > total_files:
        console.print(f"\n[green]✅ Successfully expanded dataset from {total_files} files to {total_segments} training segments![/green]")
    else:
        console.print(f"\n[yellow]⚠️  Few events detected. Check if audio files contain clear metal detection signals.[/yellow]")


# For main segmentation function
def main():
    """
    CLI script to segment dataset using MetalDetectionEventSegmenter.
    """
    parser = argparse.ArgumentParser(
        description="Segment metal detector audio files into individual detection events"
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=Path("data"),
        help="Input dataset directory (default: data)"
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data_segmented"),
        help="Output directory for segmented data (default: data_segmented)"
    )
    parser.add_argument(
        "--in-place",
        action="store_true",
        help="Add segments to existing dataset structure instead of creating new directory"
    )
    parser.add_argument(
        "--min-event-duration",
        type=float,
        default=0.5,
        help="Minimum duration for valid detection event in seconds (default: 0.5)"
    )
    parser.add_argument(
        "--energy-threshold",
        type=float,
        default=2.0,
        help="Energy threshold factor for event detection (default: 2.0)"
    )
    
    args = parser.parse_args()
    
    # Validate input directory
    if not args.input_dir.exists():
        console.print(f"[red]Error: Input directory '{args.input_dir}' does not exist![/red]")
        return
    
    # Process dataset
    process_dataset(args.input_dir, args.output_dir, args.in_place)


if __name__ == "__main__":
    main()