app.py

import os
import sys
import subprocess
import gradio as gr
import json
import yaml
import tempfile
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import time
from pathlib import Path

# VERSA paths - these are set by the Dockerfile
VERSA_ROOT = "/app/versa"
VERSA_BIN = os.path.join(VERSA_ROOT, "versa", "bin", "scorer.py")
VERSA_CONFIG_DIR = os.path.join(VERSA_ROOT, "egs")

# Data directories - also set up by the Dockerfile
DATA_DIR = "/app/data"
UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
RESULTS_DIR = os.path.join(DATA_DIR, "results")
CONFIG_DIR = os.path.join(DATA_DIR, "configs")

# Check if VERSA is installed
def check_versa_installation():
    """Check if VERSA is properly installed"""
    if not os.path.exists(VERSA_ROOT):
        return False, "VERSA directory not found."
    
    if not os.path.exists(VERSA_BIN):
        return False, "VERSA binary not found."
    
    if not os.path.exists(VERSA_CONFIG_DIR):
        return False, "VERSA configuration directory not found."
    
    # Check if the .installation_complete file exists (created by Dockerfile)
    if not os.path.exists(os.path.join(VERSA_ROOT, ".installation_complete")):
        return False, "VERSA installation indicator file not found."
    
    return True, "VERSA is properly installed."

# Check VERSA installation at startup
versa_installed, versa_status = check_versa_installation()
print(f"VERSA installation status: {versa_status}")

# Find available metric configs
def get_available_metrics():
    """Get list of available metrics from VERSA config directory"""
    metrics = []
    
    if not versa_installed:
        # If VERSA is not installed, return an empty list
        return metrics
    
    # Get all YAML files from the egs directory
    for root, _, files in os.walk(VERSA_CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                # Get relative path from VERSA_CONFIG_DIR
                rel_path = os.path.relpath(path, VERSA_CONFIG_DIR)
                metrics.append(rel_path)
    
    # Add custom configs
    for root, _, files in os.walk(CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                rel_path = f"custom/{os.path.basename(path)}"
                metrics.append(rel_path)
    
    return sorted(metrics)

# Get all available metric names
def get_available_metric_names():
    """Get list of all available metric names in VERSA"""
    metric_names = set()
    
    if not versa_installed:
        # If VERSA is not installed, return an empty list
        return []
    
    # First check the universal metrics file
    universal_metrics_yaml = os.path.join(CONFIG_DIR, "universal_metrics.yaml")
    if os.path.exists(universal_metrics_yaml):
        try:
            with open(universal_metrics_yaml, 'r') as f:
                config = yaml.safe_load(f)
                if isinstance(config, list):
                    for item in config:
                        if isinstance(item, dict) and 'name' in item:
                            metric_names.add(item['name'])
        except Exception:
            pass
    
    # Then parse all YAML files to extract additional metric names
    for root, _, files in os.walk(VERSA_CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                try:
                    with open(path, 'r') as f:
                        config = yaml.safe_load(f)
                        if isinstance(config, list):
                            for item in config:
                                if isinstance(item, dict) and 'name' in item:
                                    metric_names.add(item['name'])
                except Exception:
                    pass
    
    return sorted(list(metric_names))

# Get metric description from YAML file
def get_metric_description(metric_path):
    """Get description of a metric from its YAML file"""
    if not versa_installed:
        return "VERSA is not installed. Metric descriptions are unavailable."
    
    if metric_path.startswith("custom/"):
        # Handle custom metrics
        filename = metric_path.split("/")[1]
        full_path = os.path.join(CONFIG_DIR, filename)
    else:
        full_path = os.path.join(VERSA_CONFIG_DIR, metric_path)
    
    try:
        with open(full_path, 'r') as f:
            config = yaml.safe_load(f)
            
            # Check if the config has a description field
            if isinstance(config, dict) and 'description' in config:
                return config.get('description', 'No description available')
            
            # If it's a list of metrics, return a summary
            if isinstance(config, list):
                metric_names = []
                for item in config:
                    if isinstance(item, dict) and 'name' in item:
                        metric_names.append(item['name'])
                
                if metric_names:
                    return f"Contains metrics: {', '.join(metric_names)}"
            
            return "No description available"
    except Exception as e:
        return f"Could not load description: {str(e)}"

# Create custom metric config file
def create_custom_metric_config(selected_metrics, metric_parameters):
    """Create a custom metric configuration file from selected metrics"""
    if not versa_installed:
        return None, "VERSA is not installed. Cannot create custom metric configuration."
    
    if not selected_metrics:
        return None, "Please select at least one metric"
    
    # Load universal metrics as reference
    universal_metrics = []
    universal_metrics_yaml = os.path.join(CONFIG_DIR, "universal_metrics.yaml")
    try:
        with open(universal_metrics_yaml, 'r') as f:
            universal_metrics = yaml.safe_load(f)
    except Exception as e:
        return None, f"Error loading universal metrics: {str(e)}"
    
    # Create new metric config
    custom_metrics = []
    for metric_name in selected_metrics:
        # Find the metric in universal metrics
        for metric in universal_metrics:
            if metric.get('name') == metric_name:
                # Add the metric to custom metrics
                custom_metrics.append(metric.copy())
                break
    
    # Apply any custom parameters from metric_parameters
    if metric_parameters:
        try:
            params = yaml.safe_load(metric_parameters)
            if isinstance(params, dict):
                for metric in custom_metrics:
                    metric_name = metric.get('name')
                    if metric_name in params and isinstance(params[metric_name], dict):
                        # Update metric parameters
                        metric.update(params[metric_name])
        except Exception as e:
            return None, f"Error parsing metric parameters: {str(e)}"
    
    # Create a custom YAML file
    timestamp = int(time.time())
    custom_yaml_path = os.path.join(CONFIG_DIR, f"custom_metrics_{timestamp}.yaml")
    
    try:
        with open(custom_yaml_path, 'w') as f:
            yaml.dump(custom_metrics, f, default_flow_style=False)
        
        return f"custom/{os.path.basename(custom_yaml_path)}", f"Custom metric configuration created successfully with {len(custom_metrics)} metrics"
    except Exception as e:
        return None, f"Error creating custom metric configuration: {str(e)}"

# Load metric config file
def load_metric_config(config_path):
    """Load a metric configuration file and return its content"""
    if not versa_installed and not config_path.startswith("custom/"):
        return "VERSA is not installed. Cannot load metric configuration."
    
    if config_path.startswith("custom/"):
        # Handle custom metrics
        filename = config_path.split("/")[1]
        full_path = os.path.join(CONFIG_DIR, filename)
    else:
        full_path = os.path.join(VERSA_CONFIG_DIR, config_path)
    
    try:
        with open(full_path, 'r') as f:
            content = f.read()
        
        return content
    except Exception as e:
        return f"Error loading metric configuration: {str(e)}"

# Save uploaded YAML file
def save_uploaded_yaml(file_obj):
    """Save an uploaded YAML file to the configs directory"""
    if file_obj is None:
        return None, "No file uploaded"
    
    try:
        # Get the file name and create a new path
        filename = os.path.basename(file_obj.name)
        if not filename.endswith('.yaml'):
            filename += '.yaml'
        
        # Ensure unique filename
        timestamp = int(time.time())
        yaml_path = os.path.join(CONFIG_DIR, f"uploaded_{timestamp}_{filename}")
        
        # Copy the file
        with open(file_obj.name, 'rb') as src, open(yaml_path, 'wb') as dst:
            dst.write(src.read())
        
        # Validate YAML format
        with open(yaml_path, 'r') as f:
            yaml.safe_load(f)
        
        return f"custom/{os.path.basename(yaml_path)}", f"YAML file uploaded successfully as {os.path.basename(yaml_path)}"
    except yaml.YAMLError:
        if os.path.exists(yaml_path):
            os.remove(yaml_path)
        return None, "Invalid YAML format. Please check your file."
    except Exception as e:
        if os.path.exists(yaml_path):
            os.remove(yaml_path)
        return None, f"Error uploading YAML file: {str(e)}"

# Process audio files and run VERSA evaluation
def evaluate_audio(gt_file, pred_file, metric_config, include_timestamps=False):
    """Evaluate audio files using VERSA"""
    if not versa_installed:
        return None, "VERSA is not installed. Evaluation cannot be performed."
    
    if gt_file is None or pred_file is None:
        return None, "Please upload both ground truth and prediction audio files."
    
    # Determine the metric config path
    if metric_config.startswith("custom/"):
        # Handle custom metrics
        filename = metric_config.split("/")[1]
        metric_config_path = os.path.join(CONFIG_DIR, filename)
    else:
        metric_config_path = os.path.join(VERSA_CONFIG_DIR, metric_config)
    
    # Create temp directory for results
    with tempfile.TemporaryDirectory() as temp_dir:
        output_file = os.path.join(temp_dir, "result.json")
        
        # Build command
        cmd = [
            sys.executable, VERSA_BIN,
            "--score_config", metric_config_path,
            "--gt", gt_file,
            "--pred", pred_file,
            "--output_file", output_file
        ]
        
        if include_timestamps:
            cmd.append("--include_timestamp")
        
        # Run VERSA evaluation
        try:
            process = subprocess.run(
                cmd,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            # Read results
            if os.path.exists(output_file):
                with open(output_file, 'r') as f:
                    results = json.load(f)
                
                # Format results as DataFrame
                if results:
                    results_df = pd.DataFrame(results)
                    return results_df, json.dumps(results, indent=2)
                else:
                    return None, "Evaluation completed but no results were generated."
            else:
                return None, "Evaluation completed but no results file was generated."
        
        except subprocess.CalledProcessError as e:
            return None, f"Error running VERSA: {e.stderr}"

# Create the Gradio interface
def create_gradio_demo():
    """Create the Gradio demo interface"""
    available_metrics = get_available_metrics()
    default_metric = "speech.yaml" if "speech.yaml" in available_metrics else available_metrics[0] if available_metrics else None
    metric_names = get_available_metric_names()
    
    with gr.Blocks(title="VERSA Speech & Audio Evaluation Demo") as demo:
        gr.Markdown("# VERSA: Versatile Evaluation of Speech and Audio")
        
        # Display installation status
        with gr.Row():
            installation_status = gr.Textbox(
                value=f"VERSA Installation Status: {'Installed' if versa_installed else 'Not Installed - ' + versa_status}",
                label="Installation Status",
                interactive=False
            )
        
        if not versa_installed:
            gr.Markdown(f"""
            ## ⚠️ VERSA Not Installed
            
            VERSA does not appear to be properly installed. The Docker container may not have been set up correctly.
            
            Error: {versa_status}
            
            Please check the Docker build logs or contact the administrator.
            """)
        else:
            gr.Markdown("Upload audio files and evaluate them using VERSA metrics.")
            
            with gr.Tabs() as tabs:
                # Standard evaluation tab
                with gr.TabItem("Standard Evaluation"):
                    with gr.Row():
                        with gr.Column():
                            gt_audio = gr.Audio(label="Ground Truth Audio", type="filepath", sources=["upload", "microphone"])
                            pred_audio = gr.Audio(label="Prediction Audio", type="filepath", sources=["upload", "microphone"])
                            
                            metric_dropdown = gr.Dropdown(
                                choices=available_metrics,
                                label="Evaluation Metric Configuration",
                                value=default_metric,
                                info="Select a pre-defined or custom metric configuration"
                            )
                            
                            with gr.Accordion("Metric Configuration Details", open=False):
                                metric_description = gr.Textbox(
                                    label="Metric Description",
                                    value=get_metric_description(default_metric) if default_metric else "",
                                    interactive=False
                                )
                                
                                metric_content = gr.Code(
                                    label="Configuration Content",
                                    language="yaml",
                                    value=load_metric_config(default_metric) if default_metric else "",
                                    interactive=False
                                )
                            
                            include_timestamps = gr.Checkbox(
                                label="Include Timestamps in Results",
                                value=False
                            )
                            
                            eval_button = gr.Button("Evaluate")
                        
                        with gr.Column():
                            results_table = gr.Dataframe(label="Evaluation Results")
                            raw_json = gr.Code(language="json", label="Raw Results")
                
                # Custom metrics creation tab
                with gr.TabItem("Create Custom Metrics"):
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("### Option 1: Select from Available Metrics")
                            
                            metrics_checklist = gr.CheckboxGroup(
                                choices=metric_names,
                                label="Available Metrics",
                                info="Select the metrics you want to include in your custom configuration"
                            )
                            
                            metric_params = gr.Code(
                                label="Custom Parameters (Optional, YAML format)",
                                language="yaml",
                                placeholder="""# Example of custom parameters
# Replace with your own as needed
pysepm:
  wss_wgt_vec: [1, 2, 3]
mcd_f0:
  f0min: 50
  f0max: 600""",
                                interactive=True
                            )
                            
                            create_custom_button = gr.Button("Create Custom Configuration")
                            custom_status = gr.Textbox(label="Status", interactive=False)
                        
                        with gr.Column():
                            gr.Markdown("### Option 2: Upload Your Own YAML File")
                            
                            uploaded_yaml = gr.File(
                                label="Upload YAML Configuration",
                                file_types=[".yaml", ".yml"],
                                type="filepath"
                            )
                            
                            upload_button = gr.Button("Upload Configuration")
                            upload_status = gr.Textbox(label="Upload Status", interactive=False)
                            
                            gr.Markdown("### Generated Configuration")
                            custom_config_path = gr.Textbox(
                                label="Configuration Path",
                                interactive=False,
                                visible=False
                            )
                            
                            custom_config_content = gr.Code(
                                label="Configuration Content",
                                language="yaml",
                                interactive=False
                            )
                
                # About tab
                with gr.TabItem("About VERSA"):
                    gr.Markdown("""
                    ## VERSA: Versatile Evaluation of Speech and Audio
                    
                    VERSA is a toolkit dedicated to collecting evaluation metrics in speech and audio quality. 
                    It provides a comprehensive connection to cutting-edge evaluation techniques and is tightly integrated with ESPnet.
                    
                    With full installation, VERSA offers over 60 metrics with 700+ metric variations based on different configurations. 
                    These metrics encompass evaluations utilizing diverse external resources, including matching and non-matching 
                    reference audio, text transcriptions, and text captions.
                    
                    ### Features
                    
                    - Pythonic interface with flexible configuration
                    - Support for various audio formats and evaluation scenarios
                    - Integration with ESPnet
                    - Batch processing capabilities
                    - Customizable evaluation metrics
                    
                    ### Citation
                    
                    ```
                    @misc{shi2024versaversatileevaluationtoolkit,
                      title={VERSA: A Versatile Evaluation Toolkit for Speech, Audio, and Music},
                      author={Jiatong Shi and Hye-jin Shim and Jinchuan Tian and Siddhant Arora and Haibin Wu and Darius Petermann and Jia Qi Yip and You Zhang and Yuxun Tang and Wangyou Zhang and Dareen Safar Alharthi and Yichen Huang and Koichi Saito and Jionghao Han and Yiwen Zhao and Chris Donahue and Shinji Watanabe},
                      year={2024},
                      eprint={2412.17667},
                      archivePrefix={arXiv},
                      primaryClass={cs.SD},
                      url={https://arxiv.org/abs/2412.17667},
                    }
                    ```
                    
                    Learn more at [VERSA GitHub Repository](https://github.com/shinjiwlab/versa).
                    """)
            
            # Event handlers
            def update_metric_details(metric_path):
                return get_metric_description(metric_path), load_metric_config(metric_path)
            
            metric_dropdown.change(
                fn=update_metric_details,
                inputs=[metric_dropdown],
                outputs=[metric_description, metric_content]
            )
            
            eval_button.click(
                fn=evaluate_audio,
                inputs=[gt_audio, pred_audio, metric_dropdown, include_timestamps],
                outputs=[results_table, raw_json]
            )
            
            # Create custom metric configuration
            def create_and_update_custom_config(selected_metrics, metric_parameters):
                config_path, status = create_custom_metric_config(selected_metrics, metric_parameters)
                if config_path:
                    content = load_metric_config(config_path)
                    # Refresh the available metrics list
                    metrics_list = get_available_metrics()
                    return status, config_path, content, gr.Dropdown.update(choices=metrics_list, value=config_path)
                else:
                    return status, None, "", gr.Dropdown.update(choices=get_available_metrics())
            
            create_custom_button.click(
                fn=create_and_update_custom_config,
                inputs=[metrics_checklist, metric_params],
                outputs=[custom_status, custom_config_path, custom_config_content, metric_dropdown]
            )
            
            # Upload YAML file
            def upload_and_update_yaml(file_obj):
                config_path, status = save_uploaded_yaml(file_obj)
                if config_path:
                    content = load_metric_config(config_path)
                    # Refresh the available metrics list
                    metrics_list = get_available_metrics()
                    return status, config_path, content, gr.Dropdown.update(choices=metrics_list, value=config_path)
                else:
                    return status, None, "", gr.Dropdown.update(choices=get_available_metrics())
            
            upload_button.click(
                fn=upload_and_update_yaml,
                inputs=[uploaded_yaml],
                outputs=[upload_status, custom_config_path, custom_config_content, metric_dropdown]
            )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_gradio_demo()
    # Use 0.0.0.0 to listen on all interfaces, which is required for Docker
    demo.launch(server_name="0.0.0.0", server_port=7860)