backend/app.py

from backend.chart import mk_variations
from backend.examples import audio_examples_tab, image_examples_tab, video_examples_tab
from flask import Flask, Response, send_from_directory
from flask_cors import CORS
import os
import logging
import pandas as pd
import json
from io import StringIO
from tools import (
    get_leaderboard_filters,
    get_old_format_dataframe,
)  # Import your function
import typing as tp
import requests
from urllib.parse import unquote


logger = logging.getLogger(__name__)
if not logger.hasHandlers():
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s"))
    logger.addHandler(handler)

logger.setLevel(logging.INFO)
logger.warning("Starting the Flask app...")

app = Flask(__name__, static_folder="../frontend/dist", static_url_path="")
CORS(app)


@app.route("/")
def index():
    logger.warning("Serving index.html")
    return send_from_directory(app.static_folder, "index.html")


@app.route("/data/<path:filename>")
def data_files(filename):
    """
    Serves csv files from the data directory.
    """
    data_dir = os.path.join(os.path.dirname(__file__), "data")
    file_path = os.path.join(data_dir, filename)
    if os.path.isfile(file_path):
        df = pd.read_csv(file_path)
        logger.info(f"Processing file: {filename}")
        if filename.endswith("benchmark.csv"):
            return get_leaderboard(df)
        elif filename.endswith("attacks_variations.csv"):
            return get_chart(df)

    return "File not found", 404


@app.route("/examples/<path:type>")
def example_files(type):
    """
    Serve example files from the examples directory.
    """

    abs_path = "https://dl.fbaipublicfiles.com/omnisealbench/"

    # Switch based on the type parameter to call the appropriate tab function
    if type == "image":
        result = image_examples_tab(abs_path)
        return Response(json.dumps(result), mimetype="application/json")
    elif type == "audio":
        # Assuming you'll create these functions
        result = audio_examples_tab(abs_path)
        return Response(json.dumps(result), mimetype="application/json")
    elif type == "video":
        # Assuming you'll create these functions
        result = video_examples_tab(abs_path)
        return Response(json.dumps(result), mimetype="application/json")
    else:
        return "Invalid example type", 400


# Add a proxy endpoint to bypass CORS issues
@app.route("/proxy/<path:url>")
def proxy(url):
    """
    Proxy endpoint to fetch remote files and serve them to the frontend.
    This helps bypass CORS restrictions on remote resources.
    """
    try:
        # Decode the URL parameter
        url = unquote(url)

        # Make sure we're only proxying from trusted domains for security
        if not url.startswith("https://dl.fbaipublicfiles.com/"):
            return {"error": "Only proxying from allowed domains is permitted"}, 403

        response = requests.get(url, stream=True)

        if response.status_code != 200:
            return {"error": f"Failed to fetch from {url}"}, response.status_code

        # Create a Flask Response with the same content type as the original
        excluded_headers = [
            "content-encoding",
            "content-length",
            "transfer-encoding",
            "connection",
        ]
        headers = {
            name: value
            for name, value in response.headers.items()
            if name.lower() not in excluded_headers
        }

        # Add CORS headers
        headers["Access-Control-Allow-Origin"] = "*"

        return Response(response.content, response.status_code, headers)
    except Exception as e:
        return {"error": str(e)}, 500


def get_leaderboard(df):
    # Determine file type and handle accordingly

    # Modify the dataframe - you'll need to define first_cols and attack_scores
    first_cols = [
        "snr",
        "sisnr",
        "stoi",
        "pesq",
    ]  # Define appropriate values based on your needs
    attack_scores = [
        "bit_acc",
        "log10_p_value",
        "TPR",
        "FPR",
    ]  # Define appropriate values based on your needs
    categories = {
        "speed": "Time",
        "updownresample": "Time",
        "echo": "Time",
        "random_noise": "Amplitude",
        "lowpass_filter": "Amplitude",
        "highpass_filter": "Amplitude",
        "bandpass_filter": "Amplitude",
        "smooth": "Amplitude",
        "boost_audio": "Amplitude",
        "duck_audio": "Amplitude",
        "shush": "Amplitude",
        "pink_noise": "Amplitude",
        "aac_compression": "Compression",
        "mp3_compression": "Compression",
    }

    # This part adds on all the columns
    df = get_old_format_dataframe(df, first_cols, attack_scores)

    groups, default_selection = get_leaderboard_filters(df, categories)

    # Replace NaN values with None for JSON serialization
    df = df.fillna(value="NaN")

    # Transpose the DataFrame so each column becomes a row and column is the model
    df = df.set_index("model").T.reset_index()
    df = df.rename(columns={"index": "metric"})

    # Convert DataFrame to JSON
    result = {
        "groups": {group: list(metrics) for group, metrics in groups.items()},
        "default_selected_metrics": list(default_selection),
        "rows": df.to_dict(orient="records"),
    }

    return Response(json.dumps(result), mimetype="application/json")


def get_chart(df):
    # This function should return the chart data based on the DataFrame
    # For now, we will just return a placeholder response
    chart_data = mk_variations(
        df,
        # attacks_plot_metrics,
        # audio_attacks_with_variations,
    )

    return Response(json.dumps(chart_data), mimetype="application/json")


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True, use_reloader=True)