Create app.py

app.py

@@ -0,0 +1,278 @@
+import gradio as gr
+import zipfile
+import tempfile
+import shutil
+from pathlib import Path
+import pandas as pd
+import json
+import os
+import traceback
+import gc
+
+# Import your modules
+from engine import compute_mapss_measures
+from models import get_model_config, cleanup_all_models
+from config import DEFAULT_ALPHA
+from utils import clear_gpu_memory
+
+def process_audio_files(zip_file, model_name, layer, alpha):
+    """
+    Process uploaded ZIP file containing audio mixtures.
+    
+    Expected ZIP structure:
+    - references/: Contains N reference audio files
+    - outputs/: Contains N output audio files
+    """
+    
+    if zip_file is None:
+        return None, "Please upload a ZIP file"
+    
+    # Create temporary directory for processing
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+        
+        try:
+            # Extract ZIP file
+            extract_path = temp_path / "extracted"
+            extract_path.mkdir(exist_ok=True)
+            
+            with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+                zip_ref.extractall(extract_path)
+            
+            # Find references and outputs directories
+            refs_dir = None
+            outs_dir = None
+            
+            # Check for standard structure
+            for item in extract_path.iterdir():
+                if item.is_dir():
+                    if item.name.lower() in ['references', 'refs', 'reference']:
+                        refs_dir = item
+                    elif item.name.lower() in ['outputs', 'outs', 'output', 'separated']:
+                        outs_dir = item
+            
+            # If not found at root, check one level deeper
+            if refs_dir is None or outs_dir is None:
+                for item in extract_path.iterdir():
+                    if item.is_dir():
+                        for subitem in item.iterdir():
+                            if subitem.is_dir():
+                                if subitem.name.lower() in ['references', 'refs', 'reference']:
+                                    refs_dir = subitem
+                                elif subitem.name.lower() in ['outputs', 'outs', 'output', 'separated']:
+                                    outs_dir = subitem
+            
+            if refs_dir is None or outs_dir is None:
+                return None, "Could not find 'references' and 'outputs' directories in the ZIP file"
+            
+            # Get audio files
+            ref_files = sorted([f for f in refs_dir.glob("*.wav")])
+            out_files = sorted([f for f in outs_dir.glob("*.wav")])
+            
+            if len(ref_files) == 0:
+                return None, "No reference WAV files found"
+            if len(out_files) == 0:
+                return None, "No output WAV files found"
+            
+            # Create manifest
+            manifest = [{
+                "mixture_id": "uploaded_mixture",
+                "references": [str(f) for f in ref_files],
+                "systems": {
+                    "uploaded_system": [str(f) for f in out_files]
+                }
+            }]
+            
+            # Validate model and layer
+            allowed_models = set(get_model_config(0).keys())
+            if model_name not in allowed_models:
+                return None, f"Invalid model. Allowed: {', '.join(sorted(allowed_models))}"
+            
+            # Set default layer if needed
+            if model_name == "raw":
+                layer_final = 0
+            else:
+                model_defaults = {
+                    "wavlm": 24, "wav2vec2": 24, "hubert": 24,
+                    "wavlm_base": 12, "wav2vec2_base": 12, "hubert_base": 12,
+                    "wav2vec2_xlsr": 24, "ast": 12
+                }
+                layer_final = layer if layer is not None else model_defaults.get(model_name, 12)
+            
+            # Run experiment with compute_mapss_measures
+            results_dir = compute_mapss_measures(
+                models=[model_name],
+                mixtures=manifest,
+                layer=layer_final,
+                alpha=alpha,
+                verbose=True,
+                max_gpus=1,  # Limit to 1 GPU for HF Space
+                add_ci=False  # Disable CI for faster processing
+            )
+            
+            # Create output ZIP with results
+            output_zip = temp_path / "results.zip"
+            
+            with zipfile.ZipFile(output_zip, 'w') as zipf:
+                # Add all CSV files from results
+                results_path = Path(results_dir)
+                for csv_file in results_path.rglob("*.csv"):
+                    arcname = str(csv_file.relative_to(results_path.parent))
+                    zipf.write(csv_file, arcname)
+                
+                # Add params.json
+                params_file = results_path / "params.json"
+                if params_file.exists():
+                    zipf.write(params_file, str(params_file.relative_to(results_path.parent)))
+                
+                # Add manifest
+                manifest_file = results_path / "manifest_canonical.json"
+                if manifest_file.exists():
+                    zipf.write(manifest_file, str(manifest_file.relative_to(results_path.parent)))
+            
+            # Read the ZIP file to return
+            with open(output_zip, 'rb') as f:
+                output_data = f.read()
+            
+            # Create a proper file object for Gradio
+            output_file_path = temp_path / "download_results.zip"
+            with open(output_file_path, 'wb') as f:
+                f.write(output_data)
+            
+            return str(output_file_path), "Processing completed successfully!"
+            
+        except Exception as e:
+            error_msg = f"Error processing files: {str(e)}\n{traceback.format_exc()}"
+            return None, error_msg
+        finally:
+            # Ensure cleanup happens
+            cleanup_all_models()
+            clear_gpu_memory()
+            gc.collect()
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="MAPSS - Multi-source Audio Perceptual Separation Scores") as demo:
+        gr.Markdown("""
+        # MAPSS: Multi-source Audio Perceptual Separation Scores
+        
+        This tool evaluates audio source separation quality using Perceptual Similarity (PS) and Perceptual Matching (PM) metrics.
+        
+        ## How to use:
+        1. **Prepare your audio files**: Create a ZIP file with the following structure:
+           ```
+           your_mixture.zip
+           ├── references/       # Original clean sources
+           │   ├── speaker1.wav
+           │   ├── speaker2.wav
+           │   └── ...
+           └── outputs/         # Separated outputs from your algorithm
+               ├── separated1.wav
+               ├── separated2.wav
+               └── ...
+           ```
+        2. **Upload the ZIP file** using the file uploader below
+        3. **Select model and parameters**
+        4. **Click "Process"** to run the evaluation
+        5. **Download the results** as a ZIP file containing CSV files with PS/PM scores
+        
+        ## Models available:
+        - **raw**: Raw waveform features (no model)
+        - **wavlm**: WavLM Large model (best overall performance)
+        - **wav2vec2**: Wav2Vec2 Large model  
+        - **hubert**: HuBERT Large model
+        - **wavlm_base**: WavLM Base model (faster, good performance)
+        - **wav2vec2_base**: Wav2Vec2 Base model
+        - **hubert_base**: HuBERT Base model
+        - **wav2vec2_xlsr**: Wav2Vec2 XLSR-53 model (multilingual)
+        - **ast**: Audio Spectrogram Transformer
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                file_input = gr.File(
+                    label="Upload ZIP file with audio mixtures",
+                    file_types=[".zip"],
+                    type="filepath"
+                )
+                
+                model_dropdown = gr.Dropdown(
+                    choices=["raw", "wavlm", "wav2vec2", "hubert", 
+                            "wavlm_base", "wav2vec2_base", "hubert_base",
+                            "wav2vec2_xlsr", "ast"],
+                    value="wav2vec2_base",
+                    label="Select embedding model"
+                )
+                
+                layer_slider = gr.Slider(
+                    minimum=0,
+                    maximum=24,
+                    step=1,
+                    value=12,
+                    label="Layer (leave at default for automatic selection)"
+                )
+                
+                alpha_slider = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=DEFAULT_ALPHA,
+                    label="Diffusion maps alpha parameter"
+                )
+                
+                process_btn = gr.Button("Process Audio Files", variant="primary")
+            
+            with gr.Column():
+                output_file = gr.File(
+                    label="Download Results (ZIP)",
+                    type="filepath"
+                )
+                status_text = gr.Textbox(
+                    label="Status",
+                    lines=3,
+                    max_lines=10
+                )
+        
+        gr.Markdown("""
+        ## Output format:
+        The results ZIP will contain:
+        - `ps_scores_{model}.csv`: Perceptual Similarity scores for each speaker/source
+        - `pm_scores_{model}.csv`: Perceptual Matching scores for each speaker/source  
+        - `params.json`: Experiment parameters
+        - `manifest_canonical.json`: Processed file manifest
+        
+        ## Score interpretation:
+        - **PS (Perceptual Similarity)**: 0-1 score, higher is better. Measures how well the separated output matches the reference compared to other sources.
+        - **PM (Perceptual Matching)**: 0-1 score, higher is better. Measures robustness to audio distortions.
+        
+        ## Notes:
+        - Processing may take several minutes depending on the audio length and model
+        - Audio files are automatically resampled to 16kHz
+        - The tool automatically matches outputs to references based on correlation
+        - For best results, ensure equal number of reference and output files
+        
+        ## Citation:
+        If you use this tool in your research, please cite our paper (details coming soon).
+        """)
+        
+        # Set up the processing
+        process_btn.click(
+            fn=process_audio_files,
+            inputs=[file_input, model_dropdown, layer_slider, alpha_slider],
+            outputs=[output_file, status_text]
+        )
+        
+        # Add examples if you want
+        gr.Examples(
+            examples=[
+                # You can add example ZIP files here if you have them
+            ],
+            inputs=[file_input]
+        )
+    
+    return demo
+
+# Create and launch the app
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()