Update dataset configurations and descriptions for COCO 2014 and SA-1B; refactor file pattern generation in examples.py

backend/config.py

@@ -79,7 +79,47 @@ DATASET_CONFIGS = {
             "shush",
         ],
     },
-    "val2014_1k/image": {
+    "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,
         "first_cols": ["psnr", "ssim", "lpips", "decoder_time"],
@@ -164,14 +204,19 @@ DATASET_CONFIGS = {
 EXAMPLE_CONFIGS = {
     "audio": {
         "type": "audio",
-        "dataset_name": "voxpopuli",
+        "dataset_name": "voxpopuli_1k",
         "path": ABS_DATASET_PATH,
         "db_key": "voxpopuli",
     },
+    # "image": {
+    #     "dataset_name": "val2014_1k_v2",
+    #     "path": ABS_DATASET_PATH,
+    #     "db_key": "local_val2014",
+    # },
     "image": {
-        "dataset_name": "coco_val2014",
+        "dataset_name": "sa_1b_val_1k",
         "path": ABS_DATASET_PATH,
-        "db_key": "coco_val2014",
+        "db_key": "local_valid",
     },
     "video": {
         "dataset_name": "sav_val_full_v2",

backend/descriptions.py

@@ -226,13 +226,20 @@ DATASET_DESCRIPTIONS = {
         "github_link": "",
     },
 
-    "val2014_1k/image": {
+    "val2014_1k_v2/image": {
         "full_name": "COCO 2014 Validation Set",
         "description": "The COCO 2014 validation set is a widely used dataset for image watermarking tasks. It contains a diverse set of images with various objects and scenes.",
         "paper_link": "https://arxiv.org/abs/1405.0312",
         "github_link": "",
     },
 
+    "sa_1b_val_1k/image": {
+        "full_name": "Segment Anything 1 Billion",
+        "description": "Segment Anything 1 Billion (SA-1B) is a dataset designed for training general-purpose object segmentation models from open world images.",
+        "paper_link": "https://arxiv.org/abs/2304.02643",
+        "github_link": "",
+    },
+
     "sav_val_full_v2/video": {
         "full_name": "SA-Video Dataset",
         "description": "The SA-Video dataset is a collection of videos designed for video watermarking tasks. It includes a variety of video content suitable for testing watermarking techniques.",

backend/examples.py

@@ -95,8 +95,7 @@ def build_description(
 
 def build_infos(abs_path: Path, datatype: str, dataset_name: str, db_key: str):
 
-    def generate_file_patterns(prefixes, extensions):
-        indices = [0, 1, 3, 4, 5]
+    def generate_file_patterns(prefixes, extensions, indices):        
         return [
             f"{prefix}_{index:05d}.{ext}"
             for prefix in prefixes
@@ -108,17 +107,19 @@ def build_infos(abs_path: Path, datatype: str, dataset_name: str, db_key: str):
         quality_metrics = ["snr", "sisnr", "stoi", "pesq"]
         extensions = ["wav"]
         datatype_abbr = "audio"
-        eval_results_path = abs_path + f"{dataset_name}_1k/examples_eval_results.json"
+        indices = [0, 1, 3, 4, 5]
     elif datatype == "image":
         quality_metrics = ["psnr", "ssim", "lpips"]
         extensions = ["png"]
         datatype_abbr = "img"
-        eval_results_path = abs_path + f"{dataset_name}_1k/examples_eval_results.json"
+        indices = list(range(20))
     elif datatype == "video":
         quality_metrics = ["psnr", "ssim", "lpips", "msssim", "vmaf"]
         extensions = ["mp4"]
         datatype_abbr = "video"
-        eval_results_path = abs_path + f"{dataset_name}/examples_eval_results.json"
+        indices = [0, 1, 3, 4, 5]
+
+    eval_results_path = abs_path + f"{dataset_name}/examples_eval_results.json"
 
     # Determine if eval_results_path is a URL or local file
     if eval_results_path.startswith("http://") or eval_results_path.startswith(
@@ -146,7 +147,7 @@ def build_infos(abs_path: Path, datatype: str, dataset_name: str, db_key: str):
         f"wmd_{datatype_abbr}",
     ]
 
-    file_patterns = generate_file_patterns(prefixes, extensions)
+    file_patterns = generate_file_patterns(prefixes, extensions, indices)
     infos = {}
     for model_name in dataset.keys():
         model_infos = {}
@@ -170,16 +171,10 @@ def build_infos(abs_path: Path, datatype: str, dataset_name: str, db_key: str):
                     model_infos[attack] = []
                     continue
 
-                if datatype == "video":
-                    file_paths = [
-                        f"{abs_path}{dataset_name}/examples/{datatype}/{model_name}/{attack}/{pattern}"
-                        for pattern in file_patterns
-                    ]
-                else:
-                    file_paths = [
-                        f"{abs_path}{dataset_name}_1k/examples/{datatype}/{model_name}/{attack}/{pattern}"
-                        for pattern in file_patterns
-                    ]
+                file_paths = [
+                    f"{abs_path}{dataset_name}/examples/{datatype}/{model_name}/{attack}/{pattern}"
+                    for pattern in file_patterns
+                ]
 
                 all_files = []