Merge pull request #130 from BitMind-AI/data/multiclass-taxonomy

data: v23 3-class image / 4-class video taxonomy with audited labels

Author: dylanuys

src/gasbench/benchmarks/common.py

@@ -223,6 +223,9 @@ def build_plan(
             "semisynthetic": len(
                 [d for d in available_datasets if d.media_type == "semisynthetic"]
             ),
+            "rendered": len(
+                [d for d in available_datasets if d.media_type == "rendered"]
+            ),
         },
     )
     logger.info(f"Sampling configuration: {json.dumps(sampling_summary.__dict__)}")

src/gasbench/benchmarks/utils/inference.py

@@ -81,37 +81,30 @@ def _create_onnx_session(model_path: str, model_type: str):
 
 
 def process_model_output(logits: np.ndarray) -> Tuple[int, np.ndarray]:
-    """
-    Process model output logits into a binary prediction and probabilities.
-    
-    Handles 3-class, 2-class, and single-output model formats.
-    For 3-class models, synthetic and semisynthetic are conflated into AI-generated.
-    
+    """Process model output logits into a class prediction and probabilities.
+
+    Handles 4-class, 3-class, 2-class, and single-output model formats.
+    The returned class index is argmax (0=real for all multimodal heads).
+    Binary SN34 collapse (real vs not-real) happens in Metrics.update.
+
     Args:
         logits: Raw model output logits (1D array from single sample)
-    
+
     Returns:
-        Tuple of (binary_prediction, probabilities)
-        - binary_prediction: 0=real, 1=AI-generated
-        - probabilities: softmax probabilities as 1D array (length 1,2, or 3)
+        Tuple of (predicted_class, probabilities)
+        - predicted_class: argmax class, or 0/1 for a single-logit head
+        - probabilities: sigmoid (len 1) or softmax (len 2/3/4)
     """
     logits = np.atleast_1d(logits).flatten()
 
     if len(logits) == 1:
         p = 1.0 / (1.0 + np.exp(-logits[0]))
         pred_probs = np.array([p], dtype=np.float64)
-        predicted_binary = int(p > 0.5)
-        return predicted_binary, pred_probs
+        predicted_class = int(p > 0.5)
+        return predicted_class, pred_probs
 
     exp_x = np.exp(logits - np.max(logits))
     pred_probs = exp_x / np.sum(exp_x)
-
-    if len(pred_probs) == 3:
-        predicted_binary = 0 if np.argmax(pred_probs) == 0 else 1
-    elif len(pred_probs) == 2:
-        predicted_binary = int(np.argmax(pred_probs))
-    else:
-        predicted_binary = int(pred_probs[0] > 0.5)
-
-    return predicted_binary, pred_probs
+    predicted_class = int(np.argmax(pred_probs))
+    return predicted_class, pred_probs
 

src/gasbench/benchmarks/utils/metrics.py

@@ -30,23 +30,25 @@ def update(
         of the same sample would. Default weight=1.0 reproduces unweighted
         behavior.
         """
-        if label == 1 and pred == 1:
+        # SN34 stays real vs not-real. Multiclass labels/preds (semi=2, rendered=3)
+        # collapse onto the same binary confusion matrix.
+        binary_label = 0 if label == 0 else 1
+        binary_pred = 0 if pred == 0 else 1
+
+        if binary_label == 1 and binary_pred == 1:
             self.true_positives += weight
-        elif label == 0 and pred == 0:
+        elif binary_label == 0 and binary_pred == 0:
             self.true_negatives += weight
-        elif label == 0 and pred == 1:
+        elif binary_label == 0 and binary_pred == 1:
             self.false_positives += weight
-        elif label == 1 and pred == 0:
+        elif binary_label == 1 and binary_pred == 0:
             self.false_negatives += weight
 
         if pred_probs is not None:
-            self.binary_y_true.append(label)
+            self.binary_y_true.append(binary_label)
             self.binary_weights.append(weight)
-            if len(pred_probs) == 3:
-                binary_prob = 1.0 - pred_probs[0]
-                self.binary_probs.append(binary_prob)
-            elif len(pred_probs) == 2:
-                self.binary_probs.append(pred_probs[1])
+            if len(pred_probs) >= 2:
+                self.binary_probs.append(float(1.0 - pred_probs[0]))
             else:
                 self.binary_probs.append(pred_probs[0])
 

src/gasbench/benchmarks/utils/pytorch_session.py

@@ -182,9 +182,9 @@ def run(self, output_names: Any, input_dict: dict) -> List[np.ndarray]:
                 )
 
         last_dim = output.shape[-1]
-        if last_dim not in (1, 2, 3) and last_dim != self._num_classes:
+        if last_dim not in (1, 2, 3, 4) and last_dim != self._num_classes:
             raise ValueError(
-                f"Model output last dimension is {last_dim}, expected 1, 2, 3, "
+                f"Model output last dimension is {last_dim}, expected 1, 2, 3, 4, "
                 f"or num_classes ({self._num_classes}). Check that model.py returns "
                 f"logits and that model_config.yaml declares the correct num_classes."
             )

src/gasbench/constants.py

@@ -1,15 +1,68 @@
 """Constants used throughout the gasbench package."""
 
-# Media type to label mapping for binary classification
-# 0 = real, 1 = AI-generated (synthetic or semisynthetic)
-# Note: Dataset configs still distinguish synthetic vs semisynthetic for provenance tracking
+# Per-modality class indices. Class 0 is always real so binary collapse is 1 - p[0].
+#
+# Image is 3-class (no rendered): 0=real, 1=synthetic, 2=semisynthetic.
+# Video is 4-class:                0=real, 1=synthetic, 2=semisynthetic, 3=rendered.
+# Audio stays binary:              0=real, 1=synthetic (semisynthetic collapsed).
+#
+# Semisynthetic media must derive materially from captured source media. An edit
+# of exclusively synthetic or rendered inputs retains that fully synthetic
+# provenance; the operation name alone (i2i, inpaint, etc.) does not determine
+# the class. The captured contribution may be spatial, temporal, or global.
+#
+# Image has no rendered class: CGI/game-engine stills are excluded from image configs.
+IMAGE_MEDIA_TYPE_TO_LABEL = {
+    "real": 0,
+    "synthetic": 1,
+    "semisynthetic": 2,
+}
+VIDEO_MEDIA_TYPE_TO_LABEL = {
+    "real": 0,
+    "synthetic": 1,
+    "semisynthetic": 2,
+    "rendered": 3,
+}
+AUDIO_MEDIA_TYPE_TO_LABEL = {
+    "real": 0,
+    "synthetic": 1,
+    "semisynthetic": 1,
+}
+
+VALID_MEDIA_TYPES = {
+    "image": frozenset(IMAGE_MEDIA_TYPE_TO_LABEL),
+    "video": frozenset(VIDEO_MEDIA_TYPE_TO_LABEL),
+    "audio": frozenset({"real", "synthetic", "semisynthetic"}),
+}
+
+_MODALITY_LABELS = {
+    "image": IMAGE_MEDIA_TYPE_TO_LABEL,
+    "video": VIDEO_MEDIA_TYPE_TO_LABEL,
+    "audio": AUDIO_MEDIA_TYPE_TO_LABEL,
+}
+
+# Legacy binary map (real vs not-real). Prefer media_type_to_label(media_type, modality).
 MEDIA_TYPE_TO_LABEL = {
     "real": 0,
     "synthetic": 1,
     "semisynthetic": 1,
 }
 
+
+def media_type_to_label(media_type: str, modality: str) -> int:
+    """Map a dataset media_type to the integer class index for this modality."""
+    table = _MODALITY_LABELS.get(modality)
+    if table is None:
+        raise KeyError(f"Unknown modality '{modality}'")
+    try:
+        return table[media_type]
+    except KeyError:
+        raise KeyError(
+            f"Invalid media_type '{media_type}' for modality '{modality}'. "
+            f"Valid: {sorted(table)}"
+        ) from None
+
+
 # Video evaluation limits
 # Caps num_frames to prevent submitted models from overwhelming the eval system.
 MAX_VIDEO_NUM_FRAMES = 64
-

src/gasbench/dataset/config.py

@@ -7,6 +7,7 @@
 import hashlib
 
 from ..logger import get_logger
+from ..constants import VALID_MEDIA_TYPES
 
 logger = get_logger(__name__)
 
@@ -46,7 +47,7 @@ class BenchmarkDatasetConfig:
     name: str
     path: str
     modality: str  # "image", "video", or "audio"
-    media_type: str  # "real", "synthetic", or "semisynthetic"
+    media_type: str  # image: real/synthetic/semisynthetic; video also allows rendered
 
     # Download parameters
     media_per_archive: int = 100
@@ -324,11 +325,14 @@ def validate_dataset_config(
             )
 
     if "media_type" in config_dict:
-        valid_media_types = ["real", "synthetic", "semisynthetic"]
+        modality = config_dict.get("modality", "image")
+        valid_media_types = VALID_MEDIA_TYPES.get(
+            modality, VALID_MEDIA_TYPES["image"]
+        )
         if config_dict["media_type"] not in valid_media_types:
             errors.append(
-                f"Dataset '{dataset_name}': Invalid media_type '{config_dict['media_type']}'. "
-                f"Must be one of {valid_media_types}"
+                f"Dataset '{dataset_name}': Invalid media_type '{config_dict['media_type']}' "
+                f"for modality '{modality}'. Must be one of {sorted(valid_media_types)}"
             )
 
     if "source" in config_dict:
@@ -609,7 +613,8 @@ def _load_modality_config(modality: str, custom_path: Optional[str] = None) -> l
             data = yaml.safe_load(f)
         return data.get("datasets", [])
 
-    # Load from split real + synthetic configs and merge
+    # Load from split real + synthetic configs and merge.
+    # legacy_images.yaml / legacy_videos.yaml are bundled but not loaded here.
     SPLIT_CONFIGS = {
         "image": ("real_images.yaml", "synthetic_images.yaml"),
         "video": ("real_videos.yaml", "synthetic_videos.yaml"),
@@ -742,12 +747,13 @@ def get_benchmark_dataset_summary() -> Dict:
     summary = {
         "image": {
             "total": len(image_datasets),
-            "active": len(
-                [d for d in image_datasets if d.media_type in ["real", "synthetic"]]
-            ),
+            "active": len(image_datasets),
             "synthetic": len(
                 [d for d in image_datasets if d.media_type == "synthetic"]
             ),
+            "semisynthetic": len(
+                [d for d in image_datasets if d.media_type == "semisynthetic"]
+            ),
             "real": len([d for d in image_datasets if d.media_type == "real"]),
             "datasets": [
                 {"name": d.name, "path": d.path, "media_type": d.media_type}
@@ -756,12 +762,16 @@ def get_benchmark_dataset_summary() -> Dict:
         },
         "video": {
             "total": len(video_datasets),
-            "active": len(
-                [d for d in video_datasets if d.media_type in ["real", "synthetic"]]
-            ),
+            "active": len(video_datasets),
             "synthetic": len(
                 [d for d in video_datasets if d.media_type == "synthetic"]
             ),
+            "semisynthetic": len(
+                [d for d in video_datasets if d.media_type == "semisynthetic"]
+            ),
+            "rendered": len(
+                [d for d in video_datasets if d.media_type == "rendered"]
+            ),
             "real": len([d for d in video_datasets if d.media_type == "real"]),
             "datasets": [
                 {"name": d.name, "path": d.path, "media_type": d.media_type}