File size: 8,322 Bytes
9eea4a2
 
6a879a9
 
cc391f1
6a879a9
b087e88
9eea4a2
 
 
 
 
b087e88
 
6a879a9
 
b5bd664
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a879a9
b5bd664
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a879a9
b5bd664
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cca4f79
6a879a9
 
 
 
 
 
 
 
b5bd664
 
6a879a9
 
 
cc391f1
 
6a879a9
 
 
 
cca4f79
 
 
6a879a9
 
 
 
 
b5bd664
 
 
 
 
 
6a879a9
 
 
 
 
 
 
b5bd664
 
b087e88
cca4f79
 
 
cc391f1
cca4f79
 
 
 
 
9eea4a2
cca4f79
 
6a879a9
cca4f79
6a879a9
 
cca4f79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cc391f1
cca4f79
 
 
 
6a879a9
 
9eea4a2
cc391f1
cca4f79
 
 
 
 
 
 
 
 
 
 
 
 
54be5f9
b5bd664
6a879a9
 
 
 
54be5f9
6a879a9
 
cc391f1
 
6a879a9
 
 
cc391f1
6a879a9
 
 
 
 
54be5f9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
# 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 pathlib import Path

ABS_DATASET_DOMAIN = "https://dl.fbaipublicfiles.com"

# Sample dataset domain and path for local loading
# Some more configuration may be required to load examples from local files.
# ABS_DATASET_DOMAIN = "./data"

ABS_DATASET_PATH = f"{ABS_DATASET_DOMAIN}/omnisealbench/"

MODALITY_CONFIG_CONSTANTS = {
    "audio": {
        "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",
        ],
    },
    "image": {
        "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",
        ],
    },
    "video": {
        "first_cols": ["psnr", "ssim", "msssim", "lpips", "vmaf", "decoder_time"],
        "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
        "categories": {
            "HorizontalFlip": "Geometric",
            "Rotate": "Geometric",
            "Resize": "Geometric",
            "Crop": "Geometric",
            "Perspective": "Geometric",
            "Brightness": "Visual",
            "Contrast": "Visual",
            "Saturation": "Visual",
            "Grayscale": "Visual",
            "Hue": "Visual",
            "JPEG": "Compression",
            "GaussianBlur": "Visual",
            "MedianFilter": "Visual",
            "H264": "Compression",
            "H264rgb": "Compression",
            "H265": "Compression",
            "VP9": "Compression",
            "H264_Crop_Brightness0": "Mixed",
            "H264_Crop_Brightness1": "Mixed",
            "H264_Crop_Brightness2": "Mixed",
            "H264_Crop_Brightness3": "Mixed",
        },
        "attacks_with_variations": [
            "Rotate",
            "Resize",
            "Crop",
            "Brightness",
            "Contrast",
            "Saturation",
            "H264",
            "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},
}


def get_user_dataset():
    datasets = defaultdict(list)
    default_local_data_dir = str(Path(__file__).parent.joinpath("data"))
    user_data_dir = os.getenv("OMNISEAL_LEADERBOARD_DATA", default_local_data_dir)
    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, dataset_name):
    """Get example configuration for a specific dataset."""
    if not dataset_name:
        raise ValueError("Dataset name is required")

    # Check if it's a valid dataset for this type
    all_datasets = get_datasets()
    if dataset_name not in all_datasets.get(type, []):
        raise ValueError(f"Unknown dataset {dataset_name} for type {type}")

    # Extract the dataset name without the type suffix
    dataset_base_name = dataset_name.split("/")[0]

    # Check if it's a user dataset
    user_datasets = get_user_dataset()
    user_data_dir = os.getenv("OMNISEAL_LEADERBOARD_DATA", "./data")

    if dataset_name in user_datasets.get(type, []):
        # It's a user dataset
        examples_config = {
            "dataset_name": dataset_base_name,
            "path": user_data_dir + "/",
            "db_key": dataset_base_name,
        }
    else:
        # It's a predefined dataset from DATASET_CONFIGS
        if dataset_name in DATASET_CONFIGS:
            dataset_config = DATASET_CONFIGS[dataset_name]
            examples_config = {
                "dataset_name": dataset_base_name,
                "path": dataset_config["path"],
                "db_key": _get_db_key_for_dataset(dataset_base_name),
            }
        else:
            raise ValueError(f"Dataset {dataset_name} not found in configurations")

    return examples_config


def _get_db_key_for_dataset(dataset_base_name):
    """Helper function to determine the database key for a dataset"""
    # Map of dataset names to their db keys
    db_key_mapping = {
        "voxpopuli_1k": "voxpopuli",
        "val2014_1k_v2": "local_val2014",
        "sa_1b_val_1k": "local_valid",
        "sav_val_full_v2": "sa-v_sav_val_videos",
        "ravdess_1k": "ravdess",  # Add mapping for ravdess dataset
    }

    return db_key_mapping.get(dataset_base_name, dataset_base_name)


def get_dataset_config(dataset_name):
    if dataset_name in DATASET_CONFIGS:
        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()
        default_local_data_dir = str(Path(__file__).parent.joinpath("data"))
        user_data_dir = os.getenv("OMNISEAL_LEADERBOARD_DATA", default_local_data_dir)
        if dataset_name in user_dataset.get(modality, []):
            cfg = {
                "type": modality,
                "path": user_data_dir,
            }
            extra_cfg = MODALITY_CONFIG_CONSTANTS.get(cfg["type"], {})
            cfg.update(extra_cfg)
            return cfg

        raise ValueError(f"Unknown dataset: {dataset_name}")