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 should be set up during the build phase
VERSA_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "versa")
VERSA_BIN = os.path.join(VERSA_ROOT, "versa", "bin", "scorer.py")
VERSA_CONFIG_DIR = os.path.join(VERSA_ROOT, "egs")

# 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. The build process may have failed."
    
    if not os.path.exists(VERSA_BIN):
        return False, "VERSA binary not found. The installation may be incomplete."
    
    if not os.path.exists(VERSA_CONFIG_DIR):
        return False, "VERSA configuration directory not found. The installation may be incomplete."
    
    # Check if the .installation_complete file exists (created by build.sh)
    if not os.path.exists(os.path.join(VERSA_ROOT, ".installation_complete")):
        return False, "VERSA installation indicator file not found. The build process may have failed."
    
    return True, "VERSA is properly installed."

# Check VERSA installation at startup
versa_installed, versa_status = check_versa_installation()
if not versa_installed:
    print(f"WARNING: {versa_status}")
    print("The application may not function correctly without VERSA.")
else:
    print("VERSA installation verified successfully.")

# Create data directory if it doesn't exist
DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "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")

for directory in [DATA_DIR, UPLOAD_DIR, RESULTS_DIR, CONFIG_DIR]:
    os.makedirs(directory, exist_ok=True)

# Save the default universal metrics YAML file
UNIVERSAL_METRICS_YAML = os.path.join(CONFIG_DIR, "universal_metrics.yaml")
if not os.path.exists(UNIVERSAL_METRICS_YAML):
    with open(UNIVERSAL_METRICS_YAML, 'w') as f:
        f.write("""# Universal Metrics Configuration for Versa
# This file contains the configuration for various universal metrics used in speech quality assessment.

# visqol metric
# -- visqol: visual quality of speech
- name: visqol
  model: default

# Word error rate with ESPnet-OWSM model
# More model_tag can be from the ESPnet huggingface https://huggingface.co/espnet .
# The default model is `espnet/owsm_v3.1_ebf`.
# --lid: the nbest language tag
- name: lid
  model_tag: default
  nbest: 1

# nomad (reference-based) metric
# -- nomad: nomad reference-based model
- name: nomad
  model_cache: versa_cache/nomad_pt-models

# srmr related metrics
# -- srmr: speech-to-reverberation modulation energy ratio
- name: srmr
  n_cochlear_filters: 23
  low_freq: 125
  min_cf: 4
  max_cf: 128
  fast: True
  norm: False

# Emotion similarity calculated based on emo2vec
# --emo2vec_similarity: the emotion similarity with emo2vec
- name: emo2vec_similarity

# noresqa related metrics
# -- noresqa: non-matching reference based speech quality assessment
- name: noresqa
  metric_type: 1 #0: NORESQA-score, 1: NORESQA-MOS

# pysepm related metrics
# -- pysepm_fwsegsnr: frequency-weighted segmental SNR
# -- pysepm_llr: Log likelihood ratio
# -- pysepm_wss: weighted spectral slope
# -- pysepm_cd: cepstral distance objective speech quality measure
# -- pysepm_Csig, pysepm_Cbak, pysepm_Covl: composite objective speech quality
# -- pysepm_csii_high, pysepm_csii_mid, pysepm_csii_low: coherence and speech intelligibility index 
# -- pysepm_ncm: normalized-covariance measure
- name: pysepm

# nisqa score for speech quality assessment
#  -- nisqa_mos_pred: NISQA MOS prediction
#  -- nisqa_noi_pred: NISQA noise prediction
#  -- nisqa_dis_pred: NISQA distortion prediction
#  -- nisqa_col_pred: NISQA color prediction
#  --nisqa_loud_pred: NISQA loudness prediction
# NOTE(jiatong): pretrain model can be downloaded with `./tools/setup_nisqa.sh`
- name: nisqa
  nisqa_model_path: ./tools/NISQA/weights/nisqa.tar

# discrete speech metrics
# -- speech_bert: speech bert score
# -- speech_bleu: speech bleu score
# -- speech_token_distance: speech token distance score
- name: discrete_speech

# mcd f0 related metrics
#  -- mcd: mel cepstral distortion
#  -- f0_corr: f0 correlation
#  -- f0_rmse: f0 root mean square error
- name: mcd_f0
  f0min: 40
  f0max: 800
  mcep_shift: 5
  mcep_fftl: 1024
  mcep_dim: 39
  mcep_alpha: 0.466
  seq_mismatch_tolerance: 0.1
  power_threshold: -20
  dtw: false

# An overall model on MOS-bench from Sheet toolkit
# --sheet_ssqa: the mos prediction from sheet_ssqa
- name: sheet_ssqa

# pesq related metrics
# -- pesq: perceptual evaluation of speech quality
- name: pesq

# stoi related metrics
# -- stoi: short-time objective intelligibility
- name: stoi

# pseudo subjective metrics
# -- utmos: UT-MOS score
# -- dnsmos: DNS-MOS score
# -- plcmos: PLC-MOS score
# -- aecmos: AEC-MOS score
- name: pseudo_mos
  predictor_types: ["utmos", "dnsmos", "plcmos", "singmos", "utmosv2"]
  predictor_args:
    utmos:
      fs: 16000
    dnsmos:
      fs: 16000
    plcmos:
      fs: 16000
    singmos:
      fs: 16000
    utmosv2:
      fs: 16000

# Word error rate with OpenAI-Whisper model
# -- whisper_wer: word error rate of openai-whisper
- name: whisper_wer
  model_tag: default
  beam_size: 1
  text_cleaner: whisper_basic

# scoreq (reference-based) metric
# -- scoreq_ref: scoreq reference-based model
- name: scoreq_ref
  data_domain: natrual
  model_cache: versa_cache/scoreq_pt-models

# scoreq (non-reference-based) metric
# -- scoreq_nr: scoreq non-reference-based model
- name: scoreq_nr
  data_domain: natural
  model_cache: versa_cache/scoreq_pt-models

# Speech Enhancement-based Metrics
# model tag can be any ESPnet-SE huggingface repo
# -- se_si_snr: the SI-SNR from a rerference speech enhancement model
- name: se_snr
  model_tag: default

# PAM: Prompting Audio-Language Models for Audio Quality Assessment
# https://github.com/soham97/PAM/tree/main

- name: pam
  repro: true
  cache_dir: versa_cache/pam
  io: soundfile
  # TEXT ENCODER CONFIG
  text_model: 'gpt2'
  text_len: 77
  transformer_embed_dim: 768
  freeze_text_encoder_weights: True
  # AUDIO ENCODER CONFIG
  audioenc_name: 'HTSAT'
  out_emb: 768
  sampling_rate: 44100
  duration: 7
  fmin: 50
  fmax: 8000 #14000 
  n_fft: 1024 # 1028 
  hop_size: 320
  mel_bins: 64
  window_size: 1024
  # PROJECTION SPACE CONFIG 
  d_proj: 1024
  temperature: 0.003
  # TRAINING AND EVALUATION CONFIG
  num_classes: 527
  batch_size: 1024
  demo: False

# Speaking rate calculating
# --speaking_rate: correct matching words/character counts
- name: speaking_rate
  model_tag: default
  beam_size: 1
  text_cleaner: whisper_basic

# Audiobox Aesthetics (Unified automatic quality assessment for speech, music, and sound.)
- name: audiobox_aesthetics
  batch_size: 1
  cache_dir: versa_cache/audiobox

# ASR-match calculating
# --asr_match_error_rate: correct matching words/character counts
- name: asr_match
  model_tag: default
  beam_size: 1
  text_cleaner: whisper_basic

# speaker related metrics
# -- spk_similarity: speaker cosine similarity
- name: speaker
  model_tag: default

# asvspoof related metrics
# -- asvspoof_score: evaluate how the generated speech is likely to be classifiied by a deepfake classifier
- name: asvspoof_score

# signal related metrics
# -- sir: signal to interference ratio
# -- sar: signal to artifact ratio
# -- sdr: signal to distortion ratio
# -- ci-sdr: scale-invariant signal to distortion ratio
# -- si-snri: scale-invariant signal to noise ratio improvement
- name: signal_metric""")

# 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
    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 = []
    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 build process may have failed.
            Please check the build logs in the Factory tab of your Hugging Face Space.
            
            Error: {versa_status}
            """)
        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]
            )