update helper functions

backend/app.py

@@ -1,33 +1,22 @@
-from backend.config import (
-    ABS_DATASET_DOMAIN,
-    get_dataset_config,
-    get_datasets,
-)
-from backend.descriptions import (
-    DATASET_DESCRIPTIONS,
-    DESCRIPTIONS,
-    METRIC_DESCRIPTIONS,
-    MODEL_DESCRIPTIONS,
-)
-from backend.examples import (
-    get_examples_tab,
-)
-from flask import Flask, Response, send_from_directory, request
-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 logging
+import mimetypes
+import os
 import typing as tp
-import requests
+from io import StringIO
 from urllib.parse import unquote
-import mimetypes
 
+import pandas as pd
+import requests
+from flask import Flask, Response, request, send_from_directory
+from flask_cors import CORS
+from tools import (get_leaderboard_filters,  # Import your function
+                   get_old_format_dataframe)
+
+from backend.config import ABS_DATASET_DOMAIN, get_dataset_config, get_datasets
+from backend.descriptions import (DATASET_DESCRIPTIONS, DESCRIPTIONS,
+                                  METRIC_DESCRIPTIONS, MODEL_DESCRIPTIONS)
+from backend.examples import get_examples_tab
 
 logger = logging.getLogger(__name__)
 if not logger.hasHandlers():

backend/config.py

@@ -1,5 +1,10 @@
 # Change these values to match your dataset structure if loading locally or from a different source.
 # IMPORTANT: When running from docker more setup is required (e.g. on Huggingface)
+import os
+from collections import defaultdict
+from copy import deepcopy
+from typing import Any, Dict
+
 ABS_DATASET_DOMAIN = "https://dl.fbaipublicfiles.com"
 
 # Sample dataset domain and path for local loading
@@ -8,10 +13,8 @@ ABS_DATASET_DOMAIN = "https://dl.fbaipublicfiles.com"
 
 ABS_DATASET_PATH = f"{ABS_DATASET_DOMAIN}/omnisealbench/"
 
-DATASET_CONFIGS = {
-    "voxpopuli_1k/audio": {
-        "type": "audio",
-        "path": ABS_DATASET_PATH,
+MODALITY_CONFIG_CONSTANTS = {
+    "audio": {
         "first_cols": [
             "snr",
             "sisnr",
@@ -49,79 +52,7 @@ DATASET_CONFIGS = {
             "shush",
         ],
     },
-    "ravdess_1k/audio": {
-        "type": "audio",
-        "path": ABS_DATASET_PATH,
-        "first_cols": ["snr", "sisnr", "stoi", "pesq"],
-        "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
-        "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",
-        },
-        "attacks_with_variations": [
-            "random_noise",
-            "lowpass_filter",
-            "highpass_filter",
-            "boost_audio",
-            "duck_audio",
-            "shush",
-        ],
-    },
-    "val2014_1k_v2/image": {
-        "type": "image",
-        "path": ABS_DATASET_PATH,
-        "first_cols": ["psnr", "ssim", "lpips", "decoder_time"],
-        "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
-        "categories": {
-            "proportion": "Geometric",
-            "collage": "Inpainting",
-            "center_crop": "Geometric",
-            "rotate": "Geometric",
-            "jpeg": "Compression",
-            "brightness": "Visual",
-            "contrast": "Visual",
-            "saturation": "Visual",
-            "sharpness": "Visual",
-            "resize": "Geometric",
-            "overlay_text": "Inpainting",
-            "hflip": "Geometric",
-            "perspective": "Geometric",
-            "median_filter": "Visual",
-            "hue": "Visual",
-            "gaussian_blur": "Visual",
-            "comb": "Mixed",
-            "avg": "Averages",
-            "none": "Baseline",
-        },
-        "attacks_with_variations": [
-            "center_crop",
-            "jpeg",
-            "brightness",
-            "contrast",
-            "saturation",
-            "sharpness",
-            "resize",
-            "perspective",
-            "median_filter",
-            "hue",
-            "gaussian_blur",
-        ],
-    },    
-    "sa_1b_val_1k/image": {
-        "type": "image",
-        "path": ABS_DATASET_PATH,
+    "image": {
         "first_cols": ["psnr", "ssim", "lpips", "decoder_time"],
         "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
         "categories": {
@@ -159,9 +90,7 @@ DATASET_CONFIGS = {
             "gaussian_blur",
         ],
     },
-    "sav_val_full_v2/video": {
-        "type": "video",
-        "path": ABS_DATASET_PATH,
+    "video": {
         "first_cols": ["psnr", "ssim", "msssim", "lpips", "vmaf", "decoder_time"],
         "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
         "categories": {
@@ -198,12 +127,20 @@ DATASET_CONFIGS = {
             "H264rgb",
             "H265",
         ],
-    },
+    }
+}
+
+DATASET_CONFIGS = {
+    "voxpopuli_1k/audio": {"type": "audio", "path": ABS_DATASET_PATH},
+    "ravdess_1k/audio": {"type": "audio", "path": ABS_DATASET_PATH},
+    "val2014_1k_v2/image": {"type": "image", "path": ABS_DATASET_PATH},
+    "sa_1b_val_1k/image": {"type": "image", "path": ABS_DATASET_PATH},
+    "sav_val_full_v2/video": {"type": "video", "path": ABS_DATASET_PATH},
 }
 
+
 EXAMPLE_CONFIGS = {
     "audio": {
-        "type": "audio",
         "dataset_name": "voxpopuli_1k",
         "path": ABS_DATASET_PATH,
         "db_key": "voxpopuli",
@@ -226,24 +163,73 @@ EXAMPLE_CONFIGS = {
 }
 
 
+def get_user_dataset():
+    datasets = defaultdict(list)
+    user_data_dir = os.getenv("OMNISEAL_LEADERBOARD_DATA", "./data")
+    if user_data_dir:
+        for user_data in os.listdir(user_data_dir):
+            if not os.path.isdir(os.path.join(user_data_dir, user_data)):
+                continue
+            user_dtype = os.listdir(os.path.join(user_data_dir, user_data, "examples"))[0]
+            datasets[user_dtype].append(user_data + "/" + user_dtype)
+
+    return datasets
+
+
 def get_datasets():
     grouped = {"audio": [], "image": [], "video": []}
     for name, cfg in DATASET_CONFIGS.items():
         dtype = cfg.get("type")
         if dtype in grouped:
             grouped[dtype].append(name)
+
+    # Add user datasets
+    user_datasets = get_user_dataset()
+    for dtype, user_names in user_datasets.items():
+        if dtype in grouped:
+            _names = [name for name in user_names if name not in grouped[dtype]]
+            grouped[dtype].extend(_names)
     return grouped
 
 
 def get_example_config(type):
-    if type in EXAMPLE_CONFIGS:
-        return EXAMPLE_CONFIGS[type]
-    else:
+    if type not in EXAMPLE_CONFIGS:
         raise ValueError(f"Unknown example type: {type}")
 
+    examples_config: Dict[str, Any] = deepcopy(EXAMPLE_CONFIGS[type])
+
+    user_datasets = get_user_dataset()
+    user_data_dir = os.getenv("OMNISEAL_LEADERBOARD_DATA", "./data")
+    if len(user_datasets) > 0:
+        assert user_data_dir, f"OMNISEAL_LEADERBOARD_DATA is reset during loading the examples for {type}. Please set it correctly"
+        for dtype, user_names in user_datasets.items():
+            if dtype == type:
+                dataset_name = user_names[0].split("/")[0]
+                path = user_data_dir + "/"
+                examples_config = {
+                    "dataset_name": dataset_name,
+                    "path": path,
+                    "db_key": dataset_name,
+                }
+    return examples_config
+
 
 def get_dataset_config(dataset_name):
     if dataset_name in DATASET_CONFIGS:
-        return DATASET_CONFIGS[dataset_name]
+        cfg = DATASET_CONFIGS[dataset_name]
+        extra_cfg = MODALITY_CONFIG_CONSTANTS.get(cfg["type"], {})
+        cfg.update(extra_cfg)
+        return cfg
     else:
+        modality = dataset_name.split("/")[-1]
+        user_dataset = get_user_dataset()
+        if dataset_name in user_dataset.get(modality, []):
+            cfg = {
+                "type": modality,
+                "path": os.getenv("OMNISEAL_LEADERBOARD_DATA"),
+            }
+            extra_cfg = MODALITY_CONFIG_CONSTANTS.get(cfg["type"], {})
+            cfg.update(extra_cfg)
+            return cfg
+
         raise ValueError(f"Unknown dataset: {dataset_name}")

backend/examples.py

@@ -1,7 +1,6 @@
 import ast
 import json
 import re
-
 from pathlib import Path
 
 import requests