app.py

import gradio as gr
import subprocess
import os
import tempfile
import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
import scipy.ndimage
from pathlib import Path
import logging
import warnings
import shutil
from typing import Tuple, Optional, Dict, Any

# Configure matplotlib for web display
plt.switch_backend('Agg')
warnings.filterwarnings('ignore')

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

class AudioAnalyzer:
    """Core class for audio analysis with modular feature extraction methods."""

    def __init__(self, temp_dir: Optional[str] = None):
        """Initialize with a temporary directory for file storage."""
        self.temp_dir = Path(temp_dir or tempfile.mkdtemp())
        self.temp_dir.mkdir(exist_ok=True)
        self.plot_files = []  # Track plot files for cleanup
        logger.info(f"Initialized temporary directory: {self.temp_dir}")

    def cleanup(self) -> None:
        """Remove temporary directory and plot files."""
        for plot_file in self.plot_files:
            if Path(plot_file).exists():
                try:
                    Path(plot_file).unlink()
                    logger.info(f"Removed plot file: {plot_file}")
                except Exception as e:
                    logger.warning(f"Failed to remove plot file {plot_file}: {str(e)}")
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir, ignore_errors=True)
            logger.info(f"Cleaned up temporary directory: {self.temp_dir}")

    def download_youtube_audio(self, video_url: str, progress=gr.Progress()) -> Tuple[Optional[str], str]:
        """Download audio from YouTube using yt-dlp."""
        if not video_url:
            return None, "Please provide a valid YouTube URL"

        progress(0.1, desc="Initializing download...")
        output_dir = self.temp_dir / "downloaded_audio"
        output_dir.mkdir(exist_ok=True)
        output_file = output_dir / "audio.mp3"

        command = [
            "yt-dlp",
            "-x",
            "--audio-format", "mp3",
            "-o", str(output_file),
            "--no-playlist",
            "--restrict-filenames",
            video_url
        ]

        try:
            progress(0.3, desc="Downloading audio...")
            subprocess.run(command, check=True, capture_output=True, text=True)
            progress(1.0, desc="Download complete!")
            return str(output_file), f"Successfully downloaded audio: {output_file.name}"
        except FileNotFoundError:
            return None, "yt-dlp not found. Install it with: pip install yt-dlp"
        except subprocess.CalledProcessError as e:
            return None, f"Download failed: {e.stderr}"
        except Exception as e:
            logger.error(f"Unexpected error during download: {str(e)}")
            return None, f"Error: {str(e)}"

    def save_plot(self, fig, filename: str) -> Optional[str]:
        """Save matplotlib figure to a temporary file and verify existence."""
        try:
            # Use NamedTemporaryFile to ensure persistence
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False, dir=self.temp_dir) as tmp_file:
                plot_path = tmp_file.name
            fig.savefig(plot_path, dpi=300, bbox_inches='tight', format='png')
            plt.close(fig)
            if not Path(plot_path).exists():
                logger.error(f"Plot file not created: {plot_path}")
                return None
            self.plot_files.append(plot_path)
            logger.info(f"Saved plot: {plot_path}")
            return str(plot_path)
        except Exception as e:
            logger.error(f"Error saving plot {filename}: {str(e)}")
            plt.close(fig)
            return None

    def extract_basic_features(self, audio_path: str, sr: int = 16000, max_duration: float = 60.0, 
                             progress=gr.Progress()) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """Extract basic audio features and generate visualizations."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)
            duration = librosa.get_duration(y=y, sr=sr)

            if duration > max_duration:
                y = y[:int(sr * max_duration)]
                duration = max_duration

            progress(0.3, desc="Computing features...")
            features: Dict[str, Any] = {
                'duration': duration,
                'sample_rate': sr,
                'samples': len(y),
                'tempo': float(librosa.beat.beat_track(y=y, sr=sr)[0]),  # Convert to float
                'mfcc': librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13),
                'spectral_centroid': librosa.feature.spectral_centroid(y=y, sr=sr)[0],
                'spectral_rolloff': librosa.feature.spectral_rolloff(y=y, sr=sr)[0],
                'zero_crossing_rate': librosa.feature.zero_crossing_rate(y)[0]
            }

            progress(0.5, desc="Computing mel spectrogram...")
            hop_length = 512
            S_mel = librosa.feature.melspectrogram(y=y, sr=sr, hop_length=hop_length, n_mels=80)
            S_dB = librosa.power_to_db(S_mel, ref=np.max)

            progress(0.8, desc="Creating visualizations...")
            fig, axes = plt.subplots(2, 2, figsize=(15, 10))

            time_axis = np.linspace(0, duration, len(y))
            axes[0, 0].plot(time_axis, y)
            axes[0, 0].set_title('Waveform')
            axes[0, 0].set_xlabel('Time (s)')
            axes[0, 0].set_ylabel('Amplitude')

            librosa.display.specshow(S_dB, sr=sr, hop_length=hop_length, 
                                    x_axis='time', y_axis='mel', ax=axes[0, 1])
            axes[0, 1].set_title('Mel Spectrogram')

            librosa.display.specshow(features['mfcc'], sr=sr, x_axis='time', ax=axes[1, 0])
            axes[1, 0].set_title('MFCC')

            times = librosa.frames_to_time(range(len(features['spectral_centroid'])), sr=sr, hop_length=hop_length)
            axes[1, 1].plot(times, features['spectral_centroid'], label='Spectral Centroid')
            axes[1, 1].plot(times, features['spectral_rolloff'], label='Spectral Rolloff')
            axes[1, 1].set_title('Spectral Features')
            axes[1, 1].set_xlabel('Time (s)')
            axes[1, 1].legend()

            plt.tight_layout()
            plot_path = self.save_plot(fig, "basic_features")
            if not plot_path:
                return None, None, "Failed to save feature visualizations"

            # Validate feature shapes
            for key in ['mfcc', 'spectral_centroid', 'spectral_rolloff', 'zero_crossing_rate']:
                if not isinstance(features[key].shape, tuple):
                    logger.error(f"Invalid shape for {key}: {features[key].shape}")
                    return None, None, f"Invalid feature shape for {key}"

            summary = f"""
**Audio Summary:**
- Duration: {duration:.2f} seconds
- Sample Rate: {sr} Hz
- Estimated Tempo: {features['tempo']:.1f} BPM
- Number of Samples: {features['samples']:,}

**Feature Shapes:**
- MFCC: {features['mfcc'].shape}
- Spectral Centroid: {features['spectral_centroid'].shape}
- Spectral Rolloff: {features['spectral_rolloff'].shape}
- Zero Crossing Rate: {features['zero_crossing_rate'].shape}
            """

            progress(1.0, desc="Analysis complete!")
            return plot_path, summary, None

        except Exception as e:
            logger.error(f"Error processing audio: {str(e)}")
            return None, None, f"Error processing audio: {str(e)}"

    def extract_chroma_features(self, audio_path: str, sr: int = 16000, max_duration: float = 30.0, 
                              progress=gr.Progress()) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """Extract and visualize enhanced chroma features."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)
            if len(y) > sr * max_duration:
                y = y[:int(sr * max_duration)]

            progress(0.3, desc="Computing chroma variants...")
            chroma_orig = librosa.feature.chroma_cqt(y=y, sr=sr)
            y_harm = librosa.effects.harmonic(y=y, margin=8)
            chroma_harm = librosa.feature.chroma_cqt(y=y_harm, sr=sr)
            chroma_filter = np.minimum(chroma_harm,
                                     librosa.decompose.nn_filter(chroma_harm,
                                                                aggregate=np.median,
                                                                metric='cosine'))
            chroma_smooth = scipy.ndimage.median_filter(chroma_filter, size=(1, 9))
            chroma_stft = librosa.feature.chroma_stft(y=y, sr=sr)
            chroma_cens = librosa.feature.chroma_cens(y=y, sr=sr)

            progress(0.8, desc="Creating visualizations...")
            fig, axes = plt.subplots(3, 2, figsize=(15, 12))
            axes = axes.flatten()

            for i, (chroma, title) in enumerate([
                (chroma_orig, 'Original Chroma (CQT)'),
                (chroma_harm, 'Harmonic Chroma'),
                (chroma_filter, 'Non-local Filtered'),
                (chroma_smooth, 'Median Filtered'),
                (chroma_stft, 'Chroma (STFT)'),
                (chroma_cens, 'CENS Features')
            ]):
                librosa.display.specshow(chroma, y_axis='chroma', x_axis='time', ax=axes[i])
                axes[i].set_title(title)

            plt.tight_layout()
            plot_path = self.save_plot(fig, "chroma_features")
            if not plot_path:
                return None, None, "Failed to save chroma visualizations"

            summary = "Chroma feature analysis complete! Visualizations show different chroma extraction methods for harmonic analysis."
            progress(1.0, desc="Chroma analysis complete!")
            return plot_path, summary, None

        except Exception as e:
            logger.error(f"Error processing chroma features: {str(e)}")
            return None, None, f"Error processing chroma features: {str(e)}"

    def generate_patches(self, audio_path: str, sr: int = 16000, patch_duration: float = 5.0, 
                        hop_duration: float = 1.0, progress=gr.Progress()) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """Generate fixed-duration patches for transformer input."""
        if not audio_path or not Path(audio_path).exists():
            return None, None, "Invalid or missing audio file"

        try:
            progress(0.1, desc="Loading audio...")
            y, sr = librosa.load(audio_path, sr=sr)

            progress(0.3, desc="Computing mel spectrogram...")
            hop_length = 512
            S_mel = librosa.feature.melspectrogram(y=y, sr=sr, hop_length=hop_length, n_mels=80)
            S_dB = librosa.power_to_db(S_mel, ref=np.max)

            progress(0.5, desc="Generating patches...")
            patch_frames = librosa.time_to_frames(patch_duration, sr=sr, hop_length=hop_length)
            hop_frames = librosa.time_to_frames(hop_duration, sr=sr, hop_length=hop_length)
            patches = librosa.util.frame(S_dB, frame_length=patch_frames, hop_length=hop_frames)

            progress(0.8, desc="Creating visualizations...")
            num_patches_to_show = min(6, patches.shape[-1])
            fig, axes = plt.subplots(2, 3, figsize=(18, 8))
            axes = axes.flatten()

            for i in range(num_patches_to_show):
                librosa.display.specshow(patches[..., i], y_axis='mel', x_axis='time', 
                                       ax=axes[i], sr=sr, hop_length=hop_length)
                axes[i].set_title(f'Patch {i+1}')

            for i in range(num_patches_to_show, len(axes)):
                axes[i].set_visible(False)

            plt.tight_layout()
            plot_path = self.save_plot(fig, "patches")
            if not plot_path:
                return None, None, "Failed to save patch visualizations"

            summary = f"""
**Patch Generation Summary:**
- Total patches generated: {patches.shape[-1]}
- Patch duration: {patch_duration:.1f} seconds
- Hop duration: {hop_duration:.1f} seconds
- Patch shape (mels, time, patches): {patches.shape}
- Each patch covers {patch_frames} time frames
            """

            progress(1.0, desc="Patch generation complete!")
            return plot_path, summary, None

        except Exception as e:
            logger.error(f"Error generating patches: {str(e)}")
            return None, None, f"Error generating patches: {str(e)}"

def create_gradio_interface() -> gr.Blocks:
    """Create a modular Gradio interface for audio analysis."""
    analyzer = AudioAnalyzer()

    with gr.Blocks(title="🎵 Audio Analysis Suite", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎵 Audio Analysis Suite

        Analyze audio from YouTube videos or uploaded files. Extract features or generate transformer patches for deep learning applications.

        **Features:**
        - 📊 **Basic Features**: Waveform, Mel Spectrogram, MFCC, Spectral Analysis, Tempo Detection
        - 🎼 **Chroma Features**: Harmonic content analysis with multiple extraction methods
        - 🧩 **Transformer Patches**: Fixed-duration patches for deep learning

        **Requirements**: Dependencies are automatically installed in Hugging Face Spaces via `requirements.txt`.
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📁 Audio Input")
                with gr.Group():
                    gr.Markdown("**Download from YouTube** (Supported formats: MP3, WAV, etc.)")
                    youtube_url = gr.Textbox(
                        label="YouTube URL",
                        placeholder="https://www.youtube.com/watch?v=...",
                    )
                    download_btn = gr.Button("📥 Download Audio", variant="primary")
                    download_status = gr.Textbox(label="Download Status", interactive=False)

                with gr.Group():
                    gr.Markdown("**Or upload audio file** (Supported formats: MP3, WAV, FLAC, etc.)")
                    audio_file = gr.Audio(
                        label="Upload Audio File",
                        type="filepath",
                    )

            with gr.Column(scale=2):
                gr.Markdown("### 🔍 Analysis Results")
                with gr.Tabs():
                    with gr.Tab("📊 Basic Features"):
                        basic_plot = gr.Image(label="Feature Visualizations")
                        basic_summary = gr.Markdown(label="Feature Summary")
                        basic_btn = gr.Button("🔍 Analyze Basic Features", variant="secondary")

                    with gr.Tab("🎼 Chroma Features"):
                        chroma_plot = gr.Image(label="Chroma Visualizations")
                        chroma_summary = gr.Markdown(label="Chroma Summary")
                        chroma_btn = gr.Button("🎼 Analyze Chroma Features", variant="secondary")

                    with gr.Tab("🧩 Transformer Patches"):
                        with gr.Row():
                            patch_duration = gr.Slider(
                                label="Patch Duration (seconds)",
                                minimum=1.0, maximum=10.0, value=5.0, step=0.5,
                            )
                            hop_duration = gr.Slider(
                                label="Hop Duration (seconds)",
                                minimum=0.1, maximum=5.0, value=1.0, step=0.1,
                            )
                        patches_plot = gr.Image(label="Generated Patches")
                        patches_summary = gr.Markdown(label="Patch Summary")
                        patches_btn = gr.Button("🧩 Generate Patches", variant="secondary")

                error_output = gr.Textbox(label="Error Messages", interactive=False)

        gr.Markdown("""
        ### ℹ️ Usage Tips
        - **Processing Limits**: 60s for basic features, 30s for chroma features for fast response
        - **YouTube Downloads**: Ensure URLs are valid and respect YouTube's terms of service
        - **Visualizations**: High-quality, suitable for research and education
        - **Storage**: Temporary files are cleaned up when the interface closes
        - **Support**: For issues, check the [GitHub repository](https://github.com/your-repo)
        """)

        # Event handlers
        download_btn.click(
            fn=analyzer.download_youtube_audio,
            inputs=[youtube_url],
            outputs=[audio_file, download_status]
        )

        basic_btn.click(
            fn=analyzer.extract_basic_features,
            inputs=[audio_file],
            outputs=[basic_plot, basic_summary, error_output]
        )

        chroma_btn.click(
            fn=analyzer.extract_chroma_features,
            inputs=[audio_file],
            outputs=[chroma_plot, chroma_summary, error_output]
        )

        patches_btn.click(
            fn=analyzer.generate_patches,
            inputs=[audio_file, patch_duration, hop_duration],
            outputs=[patches_plot, patches_summary, error_output]
        )

        audio_file.change(
            fn=analyzer.extract_basic_features,
            inputs=[audio_file],
            outputs=[basic_plot, basic_summary, error_output]
        )

        demo.unload(fn=analyzer.cleanup)

    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()