Add data pipeline scripts and configuration files for Visual Genome processing

- Introduced .gitignore to exclude unnecessary files and directories.
- Added requirements.txt for project dependencies.
- Implemented dataset.yaml with download, extraction, and processing configurations.
- Created multiple scripts for data handling: download, extract, process, and build splits.
- Developed run_baselines.py and eval_baseline.py for baseline model evaluation.
- Enhanced README.md with data pipeline instructions and expected outputs.
- Established src/data module for data utilities and processing functions.

.gitignore

@@ -0,0 +1,86 @@
+# Python cache and bytecode
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Packaging / build
+.Python
+build/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+env/
+ENV/
+
+# Test / coverage / tooling cache
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.tox/
+.nox/
+.coverage
+.coverage.*
+htmlcov/
+.hypothesis/
+
+# Notebook checkpoints
+.ipynb_checkpoints/
+
+# Logs and runtime files
+*.log
+*.out
+*.err
+
+# IDE / editor files
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Local environment files
+.env
+.env.*
+
+# ML data and artifacts
+data/raw/
+data/processed/
+checkpoints/
+runs/
+outputs/
+artifacts/
+
+.jpeg
+.jpg
+.png
+.gif
+.bmp
+.tiff
+.ico
+.webp
+.svg
+.heic
+.heif
+.heif-srgb
+.heif-srgb-alpha
+.heif-srgb-alpha-alpha

README.md

@@ -16,6 +16,8 @@ Current scope starts with `chair` and multi-label attributes like `broken`, `blu
 - `configs/train.yaml` - optimization, staging, losses, checkpointing
 - `configs/eval.yaml` - metrics, visualization, error analysis
 - `configs/experiment.yaml` - baseline/main/ablation run definitions
+- `scripts/data/` - data download, extraction, processing, split pipeline
+- `src/data/` - reusable normalization and IO utilities for data pipeline
 
 ## Config Overview
 
@@ -79,7 +81,31 @@ Use `plan/plan.md` as the source of truth for quality gates and deliverables.
 
 ## How To Run
 
-This repo currently provides config skeletons. Connect them to your training entrypoint (for example `scripts/train.py`) and evaluation entrypoint (for example `scripts/eval.py`).
+### 1) Run Chair-Only Data Pipeline (Visual Genome)
+
+Default full pipeline:
+
+`python scripts/data/run_data_pipeline.py --config configs/dataset.yaml`
+
+Run individual stages:
+
+- `python scripts/data/run_data_pipeline.py --download-only --config configs/dataset.yaml`
+- `python scripts/data/run_data_pipeline.py --extract-only --config configs/dataset.yaml`
+- `python scripts/data/run_data_pipeline.py --process-only --config configs/dataset.yaml`
+- `python scripts/data/run_data_pipeline.py --split-only --config configs/dataset.yaml`
+
+Direct scripts (if needed):
+
+- `python scripts/data/download_visual_genome.py --config configs/dataset.yaml`
+- `python scripts/data/extract_visual_genome.py --config configs/dataset.yaml`
+- `python scripts/data/process_visual_genome.py --config configs/dataset.yaml`
+- `python scripts/data/build_splits.py --config configs/dataset.yaml`
+
+Important: the processing pipeline is strict chair-only (`primary_object_class: chair`).
+
+### 2) Train/Evaluate Models
+
+This repo provides dataset + experiment configs. Connect model training to your entrypoints (for example `scripts/train.py` and `scripts/eval.py`).
 
 Typical flow:
 
@@ -90,6 +116,19 @@ Typical flow:
 
 ## Expected Outputs
 
+Data artifacts:
+
+- `data/processed/visual_genome/samples.parquet` (or `samples.csv` fallback)
+- `data/processed/visual_genome/label_vocab.json`
+- `data/processed/visual_genome/label_frequencies.json`
+- `data/processed/visual_genome/processing_report.json`
+- `data/processed/visual_genome/splits/train.json`
+- `data/processed/visual_genome/splits/val.json`
+- `data/processed/visual_genome/splits/test.json`
+- `data/processed/visual_genome/splits/split_report.json`
+
+Training/eval artifacts:
+
 - model checkpoints
 - evaluation reports (mAP/F1/per-attribute)
 - embedding analysis artifacts (UMAP/retrieval results)

configs/dataset.yaml

@@ -11,12 +11,58 @@ dataset:
   image_size: 224
   num_workers: 8
   pin_memory: true
+  expected_primary_object_class: chair
+
+download:
+  enabled: true
+  force_download: false
+  timeout_seconds: 120
+  retries: 3
+  chunk_size_bytes: 1048576
+  artifacts:
+    - name: images_metadata
+      url: https://homes.cs.washington.edu/~ranjay/visualgenome/data/dataset/image_data.json.zip
+      filename: image_data.json.zip
+      sha256: null
+    - name: objects
+      url: https://homes.cs.washington.edu/~ranjay/visualgenome/data/dataset/objects.json.zip
+      filename: objects.json.zip
+      sha256: null
+    - name: attributes
+      url: https://homes.cs.washington.edu/~ranjay/visualgenome/data/dataset/attributes.json.zip
+      filename: attributes.json.zip
+      sha256: null
+
+extraction:
+  enabled: true
+  force_extract: false
+  safe_extract: true
+  marker_file: .extract_complete
+  archives:
+    - image_data.json.zip
+    - objects.json.zip
+    - attributes.json.zip
+
+processing:
+  output_format: parquet
+  csv_fallback: true
+  processing_report_file: processing_report.json
+  label_vocab_file: label_vocab.json
+  label_frequency_file: label_frequencies.json
+  samples_file_parquet: samples.parquet
+  samples_file_csv: samples.csv
+  split_dir: splits
+  unknown_label_token: __unknown__
+  keep_unmapped_attributes: false
+  strict_chair_only: true
 
 splits:
   train: 0.8
   val: 0.1
   test: 0.1
   stratify_by_attribute_groups: true
+  multilabel_stratify: true
+  stratify_fallback: iterative_bucket
   split_seed: 42
 
 attributes:

configs/model.yaml

@@ -6,13 +6,13 @@ model:
 encoders:
   clip:
     enabled: true
-    model_name: ViT-B-32
+    model_name: vit_base_patch16_clip_224.openai
     pretrained: true
     output_dim: 512
     freeze_stage1: true
   dino:
     enabled: true
-    model_name: vit_base_patch14_dinov2
+    model_name: vit_base_patch14_dinov2.lvd142m
     pretrained: true
     output_dim: 768
     freeze_stage1: true

configs/train.yaml

@@ -4,6 +4,10 @@ run:
   experiment_name: clip_dino_gnn_bipartite
   output_dir: outputs
   save_every_n_epochs: 1
+  smoke_profile:
+    epochs: 2
+    batch_size: 8
+    eval_batch_size: 8
 
 training:
   stage1:

requirements.txt

@@ -0,0 +1,12 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+
+PyYAML
+pandas
+pyarrow
+torch==2.5.1+cu121
+torchvision==0.20.1+cu121
+torchaudio==2.5.1+cu121
+timm
+scikit-learn
+Pillow
+tqdm

scripts/data/build_splits.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import random
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Tuple
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import ensure_dir, load_yaml, write_json  # noqa: E402
+
+
+def to_json_safe(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {str(k): to_json_safe(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple, set)):
+        return [to_json_safe(v) for v in value]
+    if hasattr(value, "tolist"):
+        return to_json_safe(value.tolist())
+    if hasattr(value, "item"):
+        try:
+            return value.item()
+        except Exception:
+            pass
+    return value
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Build deterministic train/val/test split manifests for chair samples."
+    )
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "dataset.yaml",
+        help="Path to dataset config yaml.",
+    )
+    return parser.parse_args()
+
+
+def _load_samples_csv(path: Path) -> List[Dict[str, Any]]:
+    rows: List[Dict[str, Any]] = []
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        reader = csv.DictReader(handle)
+        for row in reader:
+            rows.append(
+                {
+                    "image_id": int(row["image_id"]),
+                    "object_id": int(row["object_id"]),
+                    "object_name": row.get("object_name", "chair"),
+                    "image_url": row.get("image_url", ""),
+                    "image_path": row.get("image_path", ""),
+                    "attributes_raw": json.loads(row.get("attributes_raw", "[]")),
+                    "attributes_norm": json.loads(row.get("attributes_norm", "[]")),
+                }
+            )
+    return rows
+
+
+def load_samples(processed_root: Path, processing_cfg: Dict[str, Any]) -> List[Dict[str, Any]]:
+    parquet_path = processed_root / processing_cfg.get("samples_file_parquet", "samples.parquet")
+    csv_path = processed_root / processing_cfg.get("samples_file_csv", "samples.csv")
+    if parquet_path.exists():
+        try:
+            import pandas as pd  # type: ignore
+
+            df = pd.read_parquet(parquet_path)
+            rows = df.to_dict(orient="records")
+            for row in rows:
+                row = to_json_safe(row)
+                if isinstance(row.get("attributes_raw"), str):
+                    row["attributes_raw"] = json.loads(row["attributes_raw"])
+                if isinstance(row.get("attributes_norm"), str):
+                    row["attributes_norm"] = json.loads(row["attributes_norm"])
+                row["image_id"] = int(row["image_id"])
+                row["object_id"] = int(row["object_id"])
+                row["attributes_raw"] = to_json_safe(row.get("attributes_raw", []))
+                row["attributes_norm"] = to_json_safe(row.get("attributes_norm", []))
+            return rows
+        except Exception:
+            pass
+    if csv_path.exists():
+        return _load_samples_csv(csv_path)
+    raise RuntimeError("No processed samples file found (parquet/csv). Run processing first.")
+
+
+def allocate_counts(total: int, train_ratio: float, val_ratio: float, test_ratio: float) -> Tuple[int, int, int]:
+    ratios = [train_ratio, val_ratio, test_ratio]
+    names = ["train", "val", "test"]
+    base = [int(total * r) for r in ratios]
+    remaining = total - sum(base)
+    frac = sorted(
+        [(i, (total * ratios[i]) - base[i]) for i in range(3)],
+        key=lambda x: x[1],
+        reverse=True,
+    )
+    for i in range(remaining):
+        base[frac[i % len(frac)][0]] += 1
+    counts = dict(zip(names, base))
+    return counts["train"], counts["val"], counts["test"]
+
+
+def signature(labels: Iterable[str]) -> str:
+    uniq = sorted(set(labels))
+    return "|".join(uniq) if uniq else "__no_label__"
+
+
+def digest_ids(values: Iterable[int]) -> str:
+    joined = ",".join(str(v) for v in sorted(values))
+    return hashlib.sha256(joined.encode("utf-8")).hexdigest()
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_yaml(args.config)
+    dataset_cfg = cfg["dataset"]
+    split_cfg = cfg.get("splits", {})
+    processing_cfg = cfg.get("processing", {})
+
+    processed_root = REPO_ROOT / dataset_cfg["processed_dir"]
+    ensure_dir(processed_root)
+    samples = load_samples(processed_root=processed_root, processing_cfg=processing_cfg)
+    if not samples:
+        raise RuntimeError("No samples available for split generation.")
+
+    seed = int(split_cfg.get("split_seed", 42))
+    random.seed(seed)
+
+    by_image: Dict[int, Dict[str, Any]] = {}
+    for row in samples:
+        image_id = int(row["image_id"])
+        if image_id not in by_image:
+            by_image[image_id] = {"labels": set(), "rows": []}
+        by_image[image_id]["rows"].append(row)
+        by_image[image_id]["labels"].update(row.get("attributes_norm", []))
+
+    buckets: Dict[str, List[int]] = defaultdict(list)
+    for image_id, payload in by_image.items():
+        buckets[signature(payload["labels"])].append(image_id)
+
+    train_ratio = float(split_cfg.get("train", 0.8))
+    val_ratio = float(split_cfg.get("val", 0.1))
+    test_ratio = float(split_cfg.get("test", 0.1))
+    ratio_sum = train_ratio + val_ratio + test_ratio
+    if abs(ratio_sum - 1.0) > 1e-6:
+        raise RuntimeError("Split ratios must sum to 1.0")
+
+    train_ids: List[int] = []
+    val_ids: List[int] = []
+    test_ids: List[int] = []
+
+    for _, image_ids in sorted(buckets.items(), key=lambda item: item[0]):
+        random.shuffle(image_ids)
+        n_train, n_val, n_test = allocate_counts(
+            total=len(image_ids),
+            train_ratio=train_ratio,
+            val_ratio=val_ratio,
+            test_ratio=test_ratio,
+        )
+        train_ids.extend(image_ids[:n_train])
+        val_ids.extend(image_ids[n_train : n_train + n_val])
+        test_ids.extend(image_ids[n_train + n_val : n_train + n_val + n_test])
+
+    train_set = set(train_ids)
+    val_set = set(val_ids)
+    test_set = set(test_ids)
+    if train_set & val_set or train_set & test_set or val_set & test_set:
+        raise RuntimeError("Leakage detected across train/val/test image IDs.")
+
+    split_rows = {"train": [], "val": [], "test": []}
+    for row in samples:
+        row = to_json_safe(row)
+        image_id = int(row["image_id"])
+        if image_id in train_set:
+            split_rows["train"].append(row)
+        elif image_id in val_set:
+            split_rows["val"].append(row)
+        elif image_id in test_set:
+            split_rows["test"].append(row)
+
+    split_dir = processed_root / processing_cfg.get("split_dir", "splits")
+    ensure_dir(split_dir)
+    for split_name in ("train", "val", "test"):
+        write_json(split_dir / f"{split_name}.json", {"samples": split_rows[split_name]})
+
+    summary = {
+        "seed": seed,
+        "total_rows": len(samples),
+        "total_unique_images": len(by_image),
+        "train_rows": len(split_rows["train"]),
+        "val_rows": len(split_rows["val"]),
+        "test_rows": len(split_rows["test"]),
+        "train_unique_images": len(train_set),
+        "val_unique_images": len(val_set),
+        "test_unique_images": len(test_set),
+        "train_image_digest": digest_ids(train_set),
+        "val_image_digest": digest_ids(val_set),
+        "test_image_digest": digest_ids(test_set),
+        "stratify_mode": "multilabel_signature_bucket",
+        "fallback_used": str(split_cfg.get("stratify_fallback", "iterative_bucket")),
+        "leakage_check_passed": True,
+    }
+    write_json(split_dir / "split_report.json", summary)
+    print("[ok] Split manifests created with leakage check passed.")
+
+
+if __name__ == "__main__":
+    main()

scripts/data/download_visual_genome.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import download_file, ensure_dir, load_yaml, verify_checksum  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Download Visual Genome raw artifacts.")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "dataset.yaml",
+        help="Path to dataset config yaml.",
+    )
+    parser.add_argument(
+        "--force-download",
+        action="store_true",
+        help="Force re-download even if files already exist.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_yaml(args.config)
+    dataset_cfg = cfg["dataset"]
+    download_cfg = cfg.get("download", {})
+
+    if not download_cfg.get("enabled", True):
+        print("Download disabled in config. Nothing to do.")
+        return
+
+    raw_root = REPO_ROOT / dataset_cfg["root_dir"]
+    ensure_dir(raw_root)
+
+    timeout_seconds = int(download_cfg.get("timeout_seconds", 120))
+    retries = int(download_cfg.get("retries", 3))
+    chunk_size = int(download_cfg.get("chunk_size_bytes", 1024 * 1024))
+    force_download = args.force_download or bool(download_cfg.get("force_download", False))
+
+    artifacts = download_cfg.get("artifacts", [])
+    if not artifacts:
+        raise RuntimeError("No download artifacts configured.")
+
+    for artifact in artifacts:
+        filename = artifact["filename"]
+        url = artifact["url"]
+        expected_sha256 = artifact.get("sha256")
+        target = raw_root / filename
+
+        if target.exists() and not force_download:
+            if verify_checksum(target, expected_sha256):
+                print(f"[skip] {filename} already exists and checksum passed.")
+                continue
+            print(f"[warn] {filename} exists but checksum mismatch. Re-downloading.")
+
+        print(f"[download] {filename} <- {url}")
+        download_file(
+            url=url,
+            destination=target,
+            timeout_seconds=timeout_seconds,
+            retries=retries,
+            chunk_size=chunk_size,
+        )
+
+        if not verify_checksum(target, expected_sha256):
+            raise RuntimeError(f"Checksum validation failed for {target}")
+        print(f"[ok] {filename}")
+
+    print("All configured Visual Genome artifacts are downloaded.")
+
+
+if __name__ == "__main__":
+    main()

scripts/data/extract_visual_genome.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import ensure_dir, extract_archive, load_yaml  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Extract Visual Genome raw archives.")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "dataset.yaml",
+        help="Path to dataset config yaml.",
+    )
+    parser.add_argument(
+        "--force-extract",
+        action="store_true",
+        help="Force extraction even if marker exists.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_yaml(args.config)
+    dataset_cfg = cfg["dataset"]
+    extraction_cfg = cfg.get("extraction", {})
+
+    if not extraction_cfg.get("enabled", True):
+        print("Extraction disabled in config. Nothing to do.")
+        return
+
+    raw_root = REPO_ROOT / dataset_cfg["root_dir"]
+    ensure_dir(raw_root)
+
+    marker_file = raw_root / extraction_cfg.get("marker_file", ".extract_complete")
+    force_extract = args.force_extract or bool(extraction_cfg.get("force_extract", False))
+    safe_extract = bool(extraction_cfg.get("safe_extract", True))
+
+    if marker_file.exists() and not force_extract:
+        print(f"[skip] Extraction marker found: {marker_file}")
+        return
+
+    archives = extraction_cfg.get("archives", [])
+    if not archives:
+        raise RuntimeError("No extraction archives configured.")
+
+    for archive_name in archives:
+        archive_path = raw_root / archive_name
+        if not archive_path.exists():
+            raise RuntimeError(f"Archive missing: {archive_path}")
+        print(f"[extract] {archive_name}")
+        extract_archive(archive_path=archive_path, output_dir=raw_root, safe_extract=safe_extract)
+
+    marker_file.write_text("ok\n", encoding="utf-8")
+    print(f"[ok] Extraction completed. Marker written at {marker_file}")
+
+
+if __name__ == "__main__":
+    main()

scripts/data/process_visual_genome.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Tuple
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import ensure_dir, load_yaml, write_json  # noqa: E402
+from data.normalization import (  # noqa: E402
+    build_label_vocab,
+    filter_by_min_support,
+    normalize_labels,
+    normalize_text,
+)
+from data.schemas import SampleRecord  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Process Visual Genome attributes to chair-only normalized samples."
+    )
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "dataset.yaml",
+        help="Path to dataset config yaml.",
+    )
+    return parser.parse_args()
+
+
+def flatten_attribute_groups(groups_cfg: Dict[str, Iterable[str]]) -> List[str]:
+    labels: List[str] = []
+    for values in groups_cfg.values():
+        labels.extend(values)
+    return labels
+
+
+def find_json_file(raw_root: Path, target_name: str) -> Path:
+    direct = raw_root / target_name
+    if direct.exists():
+        return direct
+    matches = list(raw_root.rglob(target_name))
+    if not matches:
+        raise RuntimeError(f"Required json file not found: {target_name}")
+    return matches[0]
+
+
+def load_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def build_image_metadata_map(image_data: List[Dict[str, Any]]) -> Dict[int, Dict[str, str]]:
+    out: Dict[int, Dict[str, str]] = {}
+    for row in image_data:
+        image_id = int(row.get("image_id"))
+        image_url = str(row.get("url") or "")
+        file_name = Path(image_url).name if image_url else ""
+        image_path = str(Path("data/raw/visual_genome/images") / file_name) if file_name else ""
+        out[image_id] = {"image_url": image_url, "image_path": image_path}
+    return out
+
+
+def as_list(value: Any) -> List[str]:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return [str(v) for v in value]
+    return [str(value)]
+
+
+def object_is_chair(names: List[str], lowercase: bool, lemmatize: bool) -> bool:
+    normalized = [
+        normalize_text(name, lowercase=lowercase, lemmatize=lemmatize) for name in names
+    ]
+    return "chair" in set(normalized)
+
+
+def parse_chair_records(
+    attributes_data: List[Dict[str, Any]],
+    image_meta: Dict[int, Dict[str, str]],
+    lowercase: bool,
+    lemmatize: bool,
+    synonym_map: Dict[str, str],
+    keep_unmapped_attributes: bool,
+    allowed_labels: List[str],
+) -> Tuple[List[SampleRecord], Dict[str, int]]:
+    allowed_set = set(
+        normalize_text(label, lowercase=lowercase, lemmatize=lemmatize) for label in allowed_labels
+    )
+    row_stats = {
+        "total_attribute_objects": 0,
+        "non_chair_objects_dropped": 0,
+        "chair_objects_kept_pre_filter": 0,
+    }
+    output: List[SampleRecord] = []
+
+    for image_item in attributes_data:
+        image_id = int(image_item.get("image_id"))
+        image_attributes = image_item.get("attributes", [])
+        metadata = image_meta.get(image_id, {})
+        for obj in image_attributes:
+            row_stats["total_attribute_objects"] += 1
+            names = as_list(obj.get("names"))
+            if not object_is_chair(names, lowercase=lowercase, lemmatize=lemmatize):
+                row_stats["non_chair_objects_dropped"] += 1
+                continue
+
+            object_id = int(obj.get("object_id") or 0)
+            raw_attributes = as_list(obj.get("attributes"))
+            normalized = normalize_labels(
+                labels=raw_attributes,
+                lowercase=lowercase,
+                lemmatize=lemmatize,
+                synonym_map=synonym_map,
+                keep_unmapped=keep_unmapped_attributes,
+                allowed_labels=allowed_set,
+            )
+            record = SampleRecord(
+                image_id=image_id,
+                object_id=object_id,
+                object_name="chair",
+                image_url=metadata.get("image_url", ""),
+                image_path=metadata.get("image_path", ""),
+                attributes_raw=raw_attributes,
+                attributes_norm=normalized,
+            )
+            output.append(record)
+            row_stats["chair_objects_kept_pre_filter"] += 1
+
+    return output, row_stats
+
+
+def deduplicate_records(records: List[SampleRecord]) -> List[SampleRecord]:
+    seen = set()
+    out: List[SampleRecord] = []
+    for record in records:
+        key = (record.image_id, record.object_id)
+        if key in seen:
+            continue
+        seen.add(key)
+        out.append(record)
+    return out
+
+
+def save_samples(
+    rows: List[Dict[str, Any]],
+    output_format: str,
+    csv_fallback: bool,
+    parquet_path: Path,
+    csv_path: Path,
+) -> str:
+    if output_format == "parquet":
+        try:
+            import pandas as pd  # type: ignore
+
+            df = pd.DataFrame(rows)
+            df.to_parquet(parquet_path, index=False)
+            return str(parquet_path)
+        except Exception as exc:  # pragma: no cover
+            if not csv_fallback:
+                raise RuntimeError(f"Failed to write parquet and csv fallback disabled: {exc}") from exc
+
+    ensure_dir(csv_path.parent)
+    fieldnames = [
+        "image_id",
+        "object_id",
+        "object_name",
+        "image_url",
+        "image_path",
+        "attributes_raw",
+        "attributes_norm",
+    ]
+    with csv_path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in rows:
+            row_copy = dict(row)
+            row_copy["attributes_raw"] = json.dumps(row_copy["attributes_raw"], ensure_ascii=True)
+            row_copy["attributes_norm"] = json.dumps(row_copy["attributes_norm"], ensure_ascii=True)
+            writer.writerow(row_copy)
+    return str(csv_path)
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_yaml(args.config)
+    dataset_cfg = cfg["dataset"]
+    processing_cfg = cfg.get("processing", {})
+    filtering_cfg = cfg.get("filtering", {})
+    attributes_cfg = cfg.get("attributes", {})
+
+    strict_chair_only = bool(processing_cfg.get("strict_chair_only", True))
+    primary_object = str(cfg["project"].get("primary_object_class", "chair")).strip().lower()
+    if strict_chair_only and primary_object != "chair":
+        raise RuntimeError("This pipeline currently supports strict chair-only processing.")
+
+    raw_root = REPO_ROOT / dataset_cfg["root_dir"]
+    processed_root = REPO_ROOT / dataset_cfg["processed_dir"]
+    ensure_dir(processed_root)
+
+    attributes_path = find_json_file(raw_root, "attributes.json")
+    image_data_path = find_json_file(raw_root, "image_data.json")
+    attributes_data = load_json(attributes_path)
+    image_data = load_json(image_data_path)
+
+    lowercase = bool(attributes_cfg.get("lowercase", True))
+    lemmatize = bool(attributes_cfg.get("lemmatize", True))
+    synonym_map = {
+        normalize_text(str(k), lowercase=lowercase, lemmatize=lemmatize): normalize_text(
+            str(v), lowercase=lowercase, lemmatize=lemmatize
+        )
+        for k, v in attributes_cfg.get("synonym_map", {}).items()
+    }
+    allowed_labels = flatten_attribute_groups(attributes_cfg.get("groups", {}))
+    keep_unmapped = bool(processing_cfg.get("keep_unmapped_attributes", False))
+
+    image_meta = build_image_metadata_map(image_data)
+    records, row_stats = parse_chair_records(
+        attributes_data=attributes_data,
+        image_meta=image_meta,
+        lowercase=lowercase,
+        lemmatize=lemmatize,
+        synonym_map=synonym_map,
+        keep_unmapped_attributes=keep_unmapped,
+        allowed_labels=allowed_labels,
+    )
+
+    if bool(filtering_cfg.get("deduplicate_images", True)):
+        records = deduplicate_records(records)
+
+    if bool(filtering_cfg.get("drop_empty_labels", True)):
+        records = [r for r in records if len(r.attributes_norm) > 0]
+
+    min_support = int(attributes_cfg.get("min_support", 1))
+    filtered_label_lists, label_freq = filter_by_min_support(
+        [r.attributes_norm for r in records],
+        min_support=min_support,
+    )
+    for idx, labels in enumerate(filtered_label_lists):
+        records[idx].attributes_norm = labels
+
+    if bool(filtering_cfg.get("drop_empty_labels", True)):
+        records = [r for r in records if len(r.attributes_norm) > 0]
+
+    max_samples_per_attribute = filtering_cfg.get("max_samples_per_attribute")
+    if max_samples_per_attribute is not None:
+        per_attr_limit = int(max_samples_per_attribute)
+        attr_counts: Dict[str, int] = {}
+        limited: List[SampleRecord] = []
+        for rec in records:
+            if not rec.attributes_norm:
+                continue
+            can_take = True
+            for label in rec.attributes_norm:
+                if attr_counts.get(label, 0) >= per_attr_limit:
+                    can_take = False
+                    break
+            if can_take:
+                limited.append(rec)
+                for label in rec.attributes_norm:
+                    attr_counts[label] = attr_counts.get(label, 0) + 1
+        records = limited
+
+    label_freq = {}
+    for rec in records:
+        for label in set(rec.attributes_norm):
+            label_freq[label] = label_freq.get(label, 0) + 1
+    label_freq = dict(sorted(label_freq.items()))
+    label_vocab = build_label_vocab(label_freq)
+
+    rows = [record.to_dict() for record in records]
+    samples_parquet = processed_root / processing_cfg.get("samples_file_parquet", "samples.parquet")
+    samples_csv = processed_root / processing_cfg.get("samples_file_csv", "samples.csv")
+    output_format = str(processing_cfg.get("output_format", "parquet")).lower()
+    csv_fallback = bool(processing_cfg.get("csv_fallback", True))
+    samples_written = save_samples(
+        rows=rows,
+        output_format=output_format,
+        csv_fallback=csv_fallback,
+        parquet_path=samples_parquet,
+        csv_path=samples_csv,
+    )
+
+    write_json(
+        processed_root / processing_cfg.get("label_vocab_file", "label_vocab.json"),
+        label_vocab,
+    )
+    write_json(
+        processed_root / processing_cfg.get("label_frequency_file", "label_frequencies.json"),
+        label_freq,
+    )
+    report = {
+        "raw_attributes_file": str(attributes_path),
+        "raw_image_data_file": str(image_data_path),
+        "samples_output": samples_written,
+        "strict_chair_only": strict_chair_only,
+        "primary_object_class": primary_object,
+        "total_samples_final": len(records),
+        "labels_kept": list(label_vocab.keys()),
+        "num_labels_kept": len(label_vocab),
+        "min_support": min_support,
+        **row_stats,
+    }
+    write_json(
+        processed_root / processing_cfg.get("processing_report_file", "processing_report.json"),
+        report,
+    )
+    print(f"[ok] Processed {len(records)} chair samples.")
+
+
+if __name__ == "__main__":
+    main()

scripts/data/run_data_pipeline.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+from typing import List
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run chair-only Visual Genome data pipeline.")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "dataset.yaml",
+        help="Path to dataset config yaml.",
+    )
+    parser.add_argument(
+        "--dataset-name",
+        type=str,
+        default="visual_genome",
+        help="Dataset selector. Only visual_genome is supported in this version.",
+    )
+    parser.add_argument("--download-only", action="store_true")
+    parser.add_argument("--extract-only", action="store_true")
+    parser.add_argument("--process-only", action="store_true")
+    parser.add_argument("--split-only", action="store_true")
+    parser.add_argument("--all", action="store_true", help="Run full pipeline (default behavior).")
+    return parser.parse_args()
+
+
+def run_step(script_path: Path, config_path: Path) -> None:
+    cmd: List[str] = [sys.executable, str(script_path), "--config", str(config_path)]
+    print(f"[run] {' '.join(cmd)}")
+    subprocess.run(cmd, check=True, cwd=REPO_ROOT)
+
+
+def main() -> None:
+    args = parse_args()
+    if args.dataset_name != "visual_genome":
+        raise RuntimeError("Only --dataset-name visual_genome is supported currently.")
+
+    selected_flags = [
+        args.download_only,
+        args.extract_only,
+        args.process_only,
+        args.split_only,
+        args.all,
+    ]
+    if sum(1 for flag in selected_flags if flag) > 1:
+        raise RuntimeError("Use only one stage selector at a time.")
+
+    scripts_dir = REPO_ROOT / "scripts" / "data"
+    steps = {
+        "download": scripts_dir / "download_visual_genome.py",
+        "extract": scripts_dir / "extract_visual_genome.py",
+        "process": scripts_dir / "process_visual_genome.py",
+        "split": scripts_dir / "build_splits.py",
+    }
+
+    if args.download_only:
+        pipeline = ["download"]
+    elif args.extract_only:
+        pipeline = ["extract"]
+    elif args.process_only:
+        pipeline = ["process"]
+    elif args.split_only:
+        pipeline = ["split"]
+    else:
+        pipeline = ["download", "extract", "process", "split"]
+
+    for step in pipeline:
+        run_step(script_path=steps[step], config_path=args.config)
+    print("[ok] Data pipeline completed.")
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_baseline.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Show saved metrics for a trained baseline run.")
+    parser.add_argument("--run-name", type=str, required=True)
+    parser.add_argument("--output-dir", type=Path, default=REPO_ROOT / "outputs")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    metrics_path = args.output_dir / args.run_name / "metrics.json"
+    if not metrics_path.exists():
+        raise RuntimeError(f"Metrics not found: {metrics_path}")
+    with metrics_path.open("r", encoding="utf-8") as handle:
+        payload: Dict[str, Any] = json.load(handle)
+    report = {
+        "run_name": payload.get("run_name"),
+        "mode": payload.get("mode"),
+        "best_val_map": payload.get("best_val_map"),
+        "test_map": payload.get("test_metrics", {}).get("map"),
+        "test_macro_f1": payload.get("test_metrics", {}).get("macro_f1"),
+        "test_micro_f1": payload.get("test_metrics", {}).get("micro_f1"),
+    }
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/run_baselines.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def log(message: str) -> None:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    print(f"[{ts}] [run_baselines] {message}", flush=True)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Baseline A/B/C experiments.")
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument(
+        "--experiment-config",
+        type=Path,
+        default=REPO_ROOT / "configs" / "experiment.yaml",
+    )
+    parser.add_argument("--dataset-config", type=Path, default=REPO_ROOT / "configs" / "dataset.yaml")
+    parser.add_argument("--model-config", type=Path, default=REPO_ROOT / "configs" / "model.yaml")
+    parser.add_argument("--train-config", type=Path, default=REPO_ROOT / "configs" / "train.yaml")
+    parser.add_argument("--eval-config", type=Path, default=REPO_ROOT / "configs" / "eval.yaml")
+    return parser.parse_args()
+
+
+def load_yaml(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        payload = yaml.safe_load(handle)
+    return payload or {}
+
+
+def run_train(
+    mode: str,
+    run_name: str,
+    overrides: Dict[str, Any],
+    dataset_config: Path,
+    model_config: Path,
+    train_config: Path,
+    eval_config: Path,
+) -> None:
+    cmd: List[str] = [
+        sys.executable,
+        str(REPO_ROOT / "scripts" / "train_baseline.py"),
+        "--run-name",
+        run_name,
+        "--mode",
+        mode,
+        "--dataset-config",
+        str(dataset_config),
+        "--model-config",
+        str(model_config),
+        "--train-config",
+        str(train_config),
+        "--eval-config",
+        str(eval_config),
+        "--overrides-json",
+        json.dumps(overrides),
+    ]
+    env = dict(**os.environ)
+    env["PYTHONUNBUFFERED"] = "1"
+    cmd.insert(1, "-u")
+    log(f"Launching run '{run_name}' in mode='{mode}'")
+    log(f"Python executable: {sys.executable}")
+    log(f"Overrides: {json.dumps(overrides)}")
+    log(f"Command: {' '.join(cmd)}")
+    started = time.time()
+    completed = subprocess.run(cmd, check=False, cwd=REPO_ROOT, env=env)
+    elapsed = time.time() - started
+    log(f"Run '{run_name}' exited with code={completed.returncode} after {elapsed:.1f}s")
+    if completed.returncode != 0:
+        raise RuntimeError(f"Baseline run failed: {run_name}")
+
+
+def main() -> None:
+    args = parse_args()
+    log("Starting baseline orchestrator")
+    log(f"Repo root: {REPO_ROOT}")
+    log(f"Mode: {args.mode}")
+    log(f"Experiment config: {args.experiment_config}")
+    experiment_cfg = load_yaml(args.experiment_config)
+    baseline_runs = experiment_cfg.get("baseline_runs", [])
+    if not baseline_runs:
+        raise RuntimeError("No baseline_runs found in experiment config.")
+    log(f"Discovered {len(baseline_runs)} baseline runs")
+
+    output_dir = REPO_ROOT / "outputs"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    comparison: Dict[str, Any] = {"mode": args.mode, "runs": []}
+
+    for row in baseline_runs:
+        run_name = f"{row['name']}_{args.mode}"
+        overrides = row.get("overrides", {})
+        log(f"Preparing run: {run_name}")
+        run_train(
+            mode=args.mode,
+            run_name=run_name,
+            overrides=overrides,
+            dataset_config=args.dataset_config,
+            model_config=args.model_config,
+            train_config=args.train_config,
+            eval_config=args.eval_config,
+        )
+        metrics_path = output_dir / run_name / "metrics.json"
+        log(f"Reading metrics from: {metrics_path}")
+        with metrics_path.open("r", encoding="utf-8") as handle:
+            metrics = json.load(handle)
+        comparison["runs"].append(
+            {
+                "run_name": run_name,
+                "test_map": metrics.get("test_metrics", {}).get("map"),
+                "test_macro_f1": metrics.get("test_metrics", {}).get("macro_f1"),
+                "test_micro_f1": metrics.get("test_metrics", {}).get("micro_f1"),
+            }
+        )
+
+    comparison_path = output_dir / f"baseline_comparison_{args.mode}.json"
+    with comparison_path.open("w", encoding="utf-8") as handle:
+        json.dump(comparison, handle, indent=2, ensure_ascii=True)
+    log(f"Wrote comparison report: {comparison_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train_baseline.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+import sys
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SRC_DIR = REPO_ROOT / "src"
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))
+
+from data.io_utils import write_json  # noqa: E402
+from train.config_utils import (  # noqa: E402
+    apply_dot_overrides,
+    deep_merge,
+    load_yaml,
+    parse_overrides_json,
+)
+from train.dataset import (  # noqa: E402
+    ChairAttributeDataset,
+    build_dataloader,
+    class_pos_weights,
+    load_json,
+    split_paths,
+)
+from train.encoders import build_backbones  # noqa: E402
+from train.losses import bce_logits_loss  # noqa: E402
+from train.metrics import compute_metrics  # noqa: E402
+from train.models import BaselineClassifier  # noqa: E402
+
+
+def log(message: str) -> None:
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    print(f"[{ts}] [train_baseline] {message}", flush=True)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Train baseline model (A/B/C).")
+    parser.add_argument("--dataset-config", type=Path, default=REPO_ROOT / "configs" / "dataset.yaml")
+    parser.add_argument("--model-config", type=Path, default=REPO_ROOT / "configs" / "model.yaml")
+    parser.add_argument("--train-config", type=Path, default=REPO_ROOT / "configs" / "train.yaml")
+    parser.add_argument("--eval-config", type=Path, default=REPO_ROOT / "configs" / "eval.yaml")
+    parser.add_argument("--run-name", type=str, required=True)
+    parser.add_argument("--overrides-json", type=str, default="")
+    parser.add_argument("--mode", type=str, default="full", choices=["smoke", "full"])
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"])
+    return parser.parse_args()
+
+
+def set_seed(seed: int, deterministic: bool) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = deterministic
+    torch.backends.cudnn.benchmark = not deterministic
+
+
+def load_combined_config(args: argparse.Namespace) -> Dict[str, Any]:
+    dataset_cfg = load_yaml(args.dataset_config)
+    model_cfg = load_yaml(args.model_config)
+    train_cfg = load_yaml(args.train_config)
+    eval_cfg = load_yaml(args.eval_config)
+    combined = deep_merge(dataset_cfg, model_cfg)
+    combined = deep_merge(combined, train_cfg)
+    combined = deep_merge(combined, eval_cfg)
+    overrides = parse_overrides_json(args.overrides_json)
+    if overrides:
+        combined = apply_dot_overrides(combined, overrides)
+    return combined
+
+
+def build_loaders(cfg: Dict[str, Any], batch_size: int, eval_batch_size: int, smoke: bool) -> Tuple[Any, Any, Any, Dict[str, int]]:
+    processed_root = REPO_ROOT / cfg["dataset"]["processed_dir"]
+    label_vocab_path = processed_root / cfg["processing"]["label_vocab_file"]
+    label_freq_path = processed_root / cfg["processing"]["label_frequency_file"]
+    cache_dir = REPO_ROOT / cfg["dataset"]["root_dir"] / "images"
+    splits = split_paths(processed_root=processed_root, split_dir_name=cfg["processing"]["split_dir"])
+
+    label_vocab = load_json(label_vocab_path)
+    label_freq = load_json(label_freq_path)
+
+    train_ds = ChairAttributeDataset(splits["train"], label_vocab_path, REPO_ROOT, cache_dir)
+    val_ds = ChairAttributeDataset(splits["val"], label_vocab_path, REPO_ROOT, cache_dir)
+    test_ds = ChairAttributeDataset(splits["test"], label_vocab_path, REPO_ROOT, cache_dir)
+
+    num_workers = int(cfg["dataset"].get("num_workers", 0))
+    if smoke:
+        num_workers = min(num_workers, 2)
+
+    train_loader = build_dataloader(
+        dataset=train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    val_loader = build_dataloader(
+        dataset=val_ds,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    test_loader = build_dataloader(
+        dataset=test_ds,
+        batch_size=eval_batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=bool(cfg["dataset"].get("pin_memory", True)),
+    )
+    return train_loader, val_loader, test_loader, {"vocab": label_vocab, "freq": label_freq}
+
+
+def build_model(cfg: Dict[str, Any], num_labels: int, device: torch.device) -> BaselineClassifier:
+    use_clip = bool(cfg["encoders"]["clip"]["enabled"])
+    use_dino = bool(cfg["encoders"]["dino"]["enabled"])
+    clip_name = str(cfg["encoders"]["clip"].get("model_name", "vit_base_patch16_clip_224.openai"))
+    dino_name = str(cfg["encoders"]["dino"].get("model_name", "vit_base_patch14_dinov2.lvd142m"))
+    if clip_name == "ViT-B-32":
+        clip_name = "vit_base_patch16_clip_224.openai"
+    if dino_name == "vit_base_patch14_dinov2":
+        dino_name = "vit_base_patch14_dinov2.lvd142m"
+
+    clip_backbone, dino_backbone, feature_dim = build_backbones(
+        use_clip=use_clip,
+        use_dino=use_dino,
+        clip_model_name=clip_name,
+        dino_model_name=dino_name,
+        clip_pretrained=bool(cfg["encoders"]["clip"].get("pretrained", True)),
+        dino_pretrained=bool(cfg["encoders"]["dino"].get("pretrained", True)),
+    )
+    hidden_dims = list(cfg["baseline_heads"]["mlp"].get("hidden_dims", [512, 256]))
+    dropout = float(cfg["baseline_heads"]["mlp"].get("dropout", 0.2))
+    model = BaselineClassifier(
+        clip_backbone=clip_backbone,
+        dino_backbone=dino_backbone,
+        feature_dim=feature_dim,
+        hidden_dims=hidden_dims,
+        dropout=dropout,
+        num_labels=num_labels,
+    )
+    model.freeze_backbones()
+    return model.to(device)
+
+
+def run_eval(model: BaselineClassifier, loader: Any, device: torch.device, pos_weight: torch.Tensor) -> Dict[str, Any]:
+    model.eval()
+    all_logits: List[np.ndarray] = []
+    all_targets: List[np.ndarray] = []
+    losses: List[float] = []
+    with torch.no_grad():
+        for batch in loader:
+            targets = batch["targets"].to(device)
+            logits = model(batch["images"], device=device)
+            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
+            losses.append(float(loss.item()))
+            all_logits.append(logits.detach().cpu().numpy())
+            all_targets.append(targets.detach().cpu().numpy())
+    logits_np = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, 1))
+    targets_np = np.concatenate(all_targets, axis=0) if all_targets else np.zeros((0, 1))
+    metrics = compute_metrics(targets=targets_np, logits=logits_np)
+    metrics["loss"] = float(np.mean(losses)) if losses else 0.0
+    return metrics
+
+
+def train() -> None:
+    args = parse_args()
+    log(f"Starting run_name={args.run_name} mode={args.mode} device_request={args.device}")
+    log(f"dataset_config={args.dataset_config}")
+    log(f"model_config={args.model_config}")
+    log(f"train_config={args.train_config}")
+    cfg = load_combined_config(args)
+    run_cfg = cfg.get("run", {})
+    seed = int(run_cfg.get("seed", 42))
+    deterministic = bool(run_cfg.get("deterministic", True))
+    set_seed(seed=seed, deterministic=deterministic)
+    log(f"Seed set to {seed}, deterministic={deterministic}")
+
+    smoke = args.mode == "smoke"
+    train_epochs = 2 if smoke else int(cfg["training"]["stage1"].get("epochs", 20))
+    batch_size = 8 if smoke else int(cfg["batching"].get("batch_size", 32))
+    eval_batch_size = 8 if smoke else int(cfg["batching"].get("eval_batch_size", 64))
+
+    cuda_available = torch.cuda.is_available()
+    selected_device = "cuda" if (args.device == "auto" and cuda_available) else args.device
+    if args.device == "auto" and not cuda_available:
+        selected_device = "cpu"
+    if selected_device == "cuda" and not cuda_available:
+        raise RuntimeError(
+            "CUDA requested but not available in torch build/environment. "
+            "Install CUDA-enabled PyTorch in the active virtualenv."
+        )
+    device = torch.device(selected_device)
+    log(
+        f"Resolved device={device.type}, torch_cuda_available={cuda_available}, "
+        f"torch_version={torch.__version__}, torch_cuda={torch.version.cuda}"
+    )
+    log("Building dataloaders from processed manifests...")
+    train_loader, val_loader, test_loader, label_info = build_loaders(
+        cfg=cfg,
+        batch_size=batch_size,
+        eval_batch_size=eval_batch_size,
+        smoke=smoke,
+    )
+    log(
+        f"Dataloader sizes -> train_batches={len(train_loader)} val_batches={len(val_loader)} "
+        f"test_batches={len(test_loader)}"
+    )
+    label_vocab = label_info["vocab"]
+    label_freq = label_info["freq"]
+    num_labels = len(label_vocab)
+    log(f"Label space size: {num_labels}")
+    log("Building encoder backbones (first run may download pretrained weights)...")
+    model = build_model(cfg=cfg, num_labels=num_labels, device=device)
+    log("Model construction complete")
+    pos_weight = class_pos_weights(label_frequencies=label_freq, label_vocab=label_vocab).to(device)
+
+    lr = float(cfg["training"]["stage1"].get("lr_head", 1e-3))
+    weight_decay = float(cfg["optimization"].get("weight_decay", 1e-4))
+    optimizer = torch.optim.AdamW(
+        [p for p in model.parameters() if p.requires_grad],
+        lr=lr,
+        weight_decay=weight_decay,
+    )
+
+    output_root = REPO_ROOT / run_cfg.get("output_dir", "outputs")
+    run_dir = output_root / args.run_name
+    run_dir.mkdir(parents=True, exist_ok=True)
+
+    history: List[Dict[str, Any]] = []
+    best_map = -1.0
+    best_path = run_dir / "best.pt"
+    log(f"Training for epochs={train_epochs}, batch_size={batch_size}, eval_batch_size={eval_batch_size}")
+    for epoch in range(1, train_epochs + 1):
+        model.train()
+        train_losses: List[float] = []
+        for batch in tqdm(train_loader, desc=f"{args.run_name} epoch {epoch}/{train_epochs}"):
+            targets = batch["targets"].to(device)
+            logits = model(batch["images"], device=device)
+            loss = bce_logits_loss(logits, targets, pos_weight=pos_weight)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            train_losses.append(float(loss.item()))
+
+        val_metrics = run_eval(model=model, loader=val_loader, device=device, pos_weight=pos_weight)
+        epoch_row = {
+            "epoch": epoch,
+            "train_loss": float(np.mean(train_losses)) if train_losses else 0.0,
+            "val_loss": val_metrics["loss"],
+            "val_map": val_metrics["map"],
+            "val_macro_f1": val_metrics["macro_f1"],
+            "val_micro_f1": val_metrics["micro_f1"],
+        }
+        history.append(epoch_row)
+        log(
+            f"Epoch {epoch}/{train_epochs} -> train_loss={epoch_row['train_loss']:.4f} "
+            f"val_map={epoch_row['val_map']:.4f} val_macro_f1={epoch_row['val_macro_f1']:.4f}"
+        )
+        if val_metrics["map"] > best_map:
+            best_map = val_metrics["map"]
+            torch.save(
+                {
+                    "model_state": model.state_dict(),
+                    "label_vocab": label_vocab,
+                    "config": cfg,
+                },
+                best_path,
+            )
+            log(f"Saved new best checkpoint at epoch {epoch}: {best_path}")
+
+    log("Running final test evaluation using best checkpoint...")
+    checkpoint = torch.load(best_path, map_location=device)
+    model.load_state_dict(checkpoint["model_state"])
+    test_metrics = run_eval(model=model, loader=test_loader, device=device, pos_weight=pos_weight)
+    payload = {
+        "run_name": args.run_name,
+        "mode": args.mode,
+        "num_labels": num_labels,
+        "best_val_map": best_map,
+        "history": history,
+        "test_metrics": test_metrics,
+    }
+    write_json(run_dir / "metrics.json", payload)
+    write_json(run_dir / "label_vocab.json", label_vocab)
+    with (run_dir / "overrides.json").open("w", encoding="utf-8") as handle:
+        json.dump(parse_overrides_json(args.overrides_json), handle, indent=2, ensure_ascii=True)
+    log(f"Training complete for {args.run_name}")
+    print(json.dumps({"run_name": args.run_name, "test_map": test_metrics["map"]}, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    train()

src/data/__init__.py

@@ -0,0 +1 @@
+"""Data pipeline utilities and processing modules."""

src/data/io_utils.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import tarfile
+import time
+import zipfile
+from pathlib import Path
+from typing import Any, Dict, Optional
+from urllib.error import HTTPError, URLError
+from urllib.request import Request, urlopen
+
+import yaml
+
+
+def load_yaml(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle)
+    return data or {}
+
+
+def ensure_dir(path: Path) -> None:
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while True:
+            block = handle.read(chunk_size)
+            if not block:
+                break
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def write_json(path: Path, payload: Dict[str, Any]) -> None:
+    ensure_dir(path.parent)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=True)
+
+
+def read_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def _download_once(
+    url: str,
+    destination: Path,
+    timeout_seconds: int,
+    chunk_size: int,
+) -> None:
+    temp_file = destination.with_suffix(destination.suffix + ".part")
+    resume_pos = temp_file.stat().st_size if temp_file.exists() else 0
+
+    headers = {}
+    if resume_pos > 0:
+        headers["Range"] = f"bytes={resume_pos}-"
+
+    request = Request(url=url, headers=headers)
+    with urlopen(request, timeout=timeout_seconds) as response:
+        status = getattr(response, "status", 200)
+        supports_resume = status == 206 and resume_pos > 0
+        mode = "ab" if supports_resume else "wb"
+        with temp_file.open(mode) as out:
+            while True:
+                chunk = response.read(chunk_size)
+                if not chunk:
+                    break
+                out.write(chunk)
+    temp_file.replace(destination)
+
+
+def download_file(
+    url: str,
+    destination: Path,
+    timeout_seconds: int = 120,
+    retries: int = 3,
+    chunk_size: int = 1024 * 1024,
+) -> None:
+    ensure_dir(destination.parent)
+    attempts = max(1, retries)
+    for attempt in range(1, attempts + 1):
+        try:
+            _download_once(
+                url=url,
+                destination=destination,
+                timeout_seconds=timeout_seconds,
+                chunk_size=chunk_size,
+            )
+            return
+        except (HTTPError, URLError, TimeoutError, OSError) as exc:
+            if attempt == attempts:
+                raise RuntimeError(
+                    f"Download failed after {attempts} attempts for {url}: {exc}"
+                ) from exc
+            sleep_s = min(10, attempt * 2)
+            time.sleep(sleep_s)
+
+
+def verify_checksum(path: Path, expected_sha256: Optional[str]) -> bool:
+    if not expected_sha256:
+        return True
+    actual = sha256_file(path)
+    return actual.lower() == expected_sha256.lower()
+
+
+def _is_within_directory(base_dir: Path, target_path: Path) -> bool:
+    try:
+        base = base_dir.resolve(strict=False)
+        target = target_path.resolve(strict=False)
+        target.relative_to(base)
+        return True
+    except ValueError:
+        return False
+
+
+def extract_archive(archive_path: Path, output_dir: Path, safe_extract: bool = True) -> None:
+    ensure_dir(output_dir)
+    lower_name = archive_path.name.lower()
+
+    if zipfile.is_zipfile(archive_path):
+        with zipfile.ZipFile(archive_path, "r") as zf:
+            members = zf.infolist()
+            if safe_extract:
+                for member in members:
+                    target = output_dir / member.filename
+                    if not _is_within_directory(output_dir, target):
+                        raise RuntimeError(
+                            f"Unsafe zip member path detected: {member.filename}"
+                        )
+            zf.extractall(output_dir)
+        return
+
+    if tarfile.is_tarfile(archive_path):
+        with tarfile.open(archive_path, "r:*") as tf:
+            members = tf.getmembers()
+            if safe_extract:
+                for member in members:
+                    target = output_dir / member.name
+                    if not _is_within_directory(output_dir, target):
+                        raise RuntimeError(
+                            f"Unsafe tar member path detected: {member.name}"
+                        )
+            tf.extractall(output_dir)
+        return
+
+    raise RuntimeError(f"Unsupported archive format: {lower_name}")

src/data/normalization.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import re
+from collections import Counter
+from typing import Dict, Iterable, List, Sequence, Set, Tuple
+
+
+_NON_ALPHA_NUM = re.compile(r"[^a-z0-9 ]+")
+_MULTI_SPACE = re.compile(r"\s+")
+
+
+def simple_lemma(token: str) -> str:
+    if token.endswith("ies") and len(token) > 3:
+        return token[:-3] + "y"
+    if token.endswith("ses") and len(token) > 3:
+        return token[:-2]
+    if token.endswith("s") and len(token) > 3:
+        return token[:-1]
+    return token
+
+
+def normalize_text(value: str, lowercase: bool = True, lemmatize: bool = True) -> str:
+    text = value.strip()
+    if lowercase:
+        text = text.lower()
+    text = _NON_ALPHA_NUM.sub(" ", text)
+    text = _MULTI_SPACE.sub(" ", text).strip()
+    if not text:
+        return ""
+    if lemmatize:
+        text = " ".join(simple_lemma(part) for part in text.split(" "))
+    return text
+
+
+def normalize_label(
+    label: str,
+    lowercase: bool,
+    lemmatize: bool,
+    synonym_map: Dict[str, str],
+) -> str:
+    normalized = normalize_text(label, lowercase=lowercase, lemmatize=lemmatize)
+    if not normalized:
+        return ""
+    if normalized in synonym_map:
+        return synonym_map[normalized]
+    return normalized
+
+
+def normalize_labels(
+    labels: Sequence[str],
+    lowercase: bool,
+    lemmatize: bool,
+    synonym_map: Dict[str, str],
+    keep_unmapped: bool,
+    allowed_labels: Set[str],
+) -> List[str]:
+    out: List[str] = []
+    for label in labels:
+        normalized = normalize_label(
+            label=label,
+            lowercase=lowercase,
+            lemmatize=lemmatize,
+            synonym_map=synonym_map,
+        )
+        if not normalized:
+            continue
+        if normalized in allowed_labels or keep_unmapped:
+            out.append(normalized)
+    return sorted(set(out))
+
+
+def filter_by_min_support(
+    normalized_label_lists: Sequence[Sequence[str]],
+    min_support: int,
+) -> Tuple[List[List[str]], Dict[str, int]]:
+    counter: Counter[str] = Counter()
+    for labels in normalized_label_lists:
+        counter.update(set(labels))
+
+    keep_labels = {name for name, count in counter.items() if count >= min_support}
+    filtered: List[List[str]] = []
+    for labels in normalized_label_lists:
+        filtered.append(sorted([label for label in set(labels) if label in keep_labels]))
+
+    filtered_counter: Counter[str] = Counter()
+    for labels in filtered:
+        filtered_counter.update(set(labels))
+    return filtered, dict(sorted(filtered_counter.items()))
+
+
+def build_label_vocab(freq_map: Dict[str, int]) -> Dict[str, int]:
+    labels = sorted(freq_map.keys())
+    return {label: idx for idx, label in enumerate(labels)}

src/data/schemas.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Dict, List
+
+
+@dataclass
+class SampleRecord:
+    image_id: int
+    object_id: int
+    object_name: str
+    image_url: str
+    image_path: str
+    attributes_raw: List[str]
+    attributes_norm: List[str]
+
+    def to_dict(self) -> Dict[str, object]:
+        return asdict(self)

src/train/__init__.py

@@ -0,0 +1 @@
+"""Training utilities for baseline models."""

src/train/config_utils.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import copy
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import yaml
+
+
+def load_yaml(path: Path) -> Dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        payload = yaml.safe_load(handle)
+    return payload or {}
+
+
+def deep_merge(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
+    merged = copy.deepcopy(base)
+    for key, value in update.items():
+        if isinstance(value, dict) and isinstance(merged.get(key), dict):
+            merged[key] = deep_merge(merged[key], value)
+        else:
+            merged[key] = copy.deepcopy(value)
+    return merged
+
+
+def set_dot_key(target: Dict[str, Any], dot_key: str, value: Any) -> None:
+    parts = dot_key.split(".")
+    current = target
+    for part in parts[:-1]:
+        if part not in current or not isinstance(current[part], dict):
+            current[part] = {}
+        current = current[part]
+    current[parts[-1]] = value
+
+
+def apply_dot_overrides(target: Dict[str, Any], overrides: Dict[str, Any]) -> Dict[str, Any]:
+    out = copy.deepcopy(target)
+    for dot_key, value in overrides.items():
+        set_dot_key(out, dot_key, value)
+    return out
+
+
+def parse_overrides_json(raw: str | None) -> Dict[str, Any]:
+    if not raw:
+        return {}
+    payload = json.loads(raw)
+    if not isinstance(payload, dict):
+        raise ValueError("Overrides JSON must be an object.")
+    return payload

src/train/dataset.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+import json
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+from urllib.error import HTTPError, URLError
+from urllib.request import Request, urlopen
+
+import torch
+from PIL import Image, UnidentifiedImageError
+from torch.utils.data import DataLoader, Dataset
+
+
+def load_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+class ChairAttributeDataset(Dataset):
+    def __init__(
+        self,
+        split_path: Path,
+        label_vocab_path: Path,
+        repo_root: Path,
+        image_cache_dir: Path,
+    ) -> None:
+        payload = load_json(split_path)
+        if "samples" not in payload:
+            raise RuntimeError(f"Invalid split file format: {split_path}")
+        self.samples: List[Dict[str, Any]] = payload["samples"]
+        self.label_vocab: Dict[str, int] = load_json(label_vocab_path)
+        self.num_labels = len(self.label_vocab)
+        self.repo_root = repo_root
+        self.image_cache_dir = image_cache_dir
+        self.image_cache_dir.mkdir(parents=True, exist_ok=True)
+        self.corrupt_image_count = 0
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def _resolve_image_path(self, row: Dict[str, Any]) -> Path:
+        local_rel = row.get("image_path", "")
+        if local_rel:
+            local = self.repo_root / local_rel
+            if local.exists():
+                return local
+        image_id = str(row.get("image_id"))
+        return self.image_cache_dir / f"{image_id}.jpg"
+
+    def _candidate_urls(self, row: Dict[str, Any]) -> List[str]:
+        primary = str(row.get("image_url") or "").strip()
+        if not primary:
+            return []
+        candidates = [primary]
+        if "/VG_100K/" in primary:
+            candidates.append(primary.replace("/VG_100K/", "/VG_100K_2/"))
+        elif "/VG_100K_2/" in primary:
+            candidates.append(primary.replace("/VG_100K_2/", "/VG_100K/"))
+        deduped: List[str] = []
+        seen = set()
+        for url in candidates:
+            if url in seen:
+                continue
+            seen.add(url)
+            deduped.append(url)
+        return deduped
+
+    def _download_if_missing(self, row: Dict[str, Any], path: Path, force: bool = False) -> None:
+        if path.exists() and not force:
+            return
+        urls = self._candidate_urls(row)
+        if not urls:
+            raise RuntimeError(f"Missing image_url for image_id={row.get('image_id')}")
+        last_error: Exception | None = None
+        for url in urls:
+            try:
+                request = Request(url, headers={"User-Agent": "Mozilla/5.0"})
+                with urlopen(request, timeout=30) as response:
+                    image_bytes = response.read()
+                # Validate bytes are actually decodable image data before caching.
+                with Image.open(BytesIO(image_bytes)) as probe:
+                    probe.verify()
+                path.parent.mkdir(parents=True, exist_ok=True)
+                with path.open("wb") as handle:
+                    handle.write(image_bytes)
+                return
+            except (HTTPError, URLError, TimeoutError, OSError, UnidentifiedImageError) as exc:
+                last_error = exc
+                continue
+        raise RuntimeError(
+            f"Failed to download valid image for image_id={row.get('image_id')} from candidate URLs."
+        ) from last_error
+
+    def _open_local_image(self, path: Path) -> Image.Image:
+        with Image.open(path) as image:
+            return image.convert("RGB")
+
+    def _load_image(self, row: Dict[str, Any]) -> Image.Image:
+        image_path = self._resolve_image_path(row)
+        self._download_if_missing(row, image_path)
+        try:
+            return self._open_local_image(image_path)
+        except (UnidentifiedImageError, OSError):
+            self.corrupt_image_count += 1
+            if image_path.exists():
+                image_path.unlink(missing_ok=True)
+            # Retry once by re-downloading. Some VG image URLs occasionally produce
+            # partial/bad files; this makes loading resilient instead of crashing workers.
+            try:
+                self._download_if_missing(row, image_path, force=True)
+                return self._open_local_image(image_path)
+            except Exception:
+                return Image.new("RGB", (224, 224), color=(0, 0, 0))
+
+    def _labels_to_multihot(self, labels: List[str]) -> torch.Tensor:
+        target = torch.zeros(self.num_labels, dtype=torch.float32)
+        for label in labels:
+            if label in self.label_vocab:
+                target[self.label_vocab[label]] = 1.0
+        return target
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        row = self.samples[idx]
+        image = self._load_image(row)
+        labels = row.get("attributes_norm", [])
+        target = self._labels_to_multihot(labels)
+        return {
+            "image": image,
+            "target": target,
+            "image_id": int(row["image_id"]),
+            "labels": labels,
+        }
+
+
+def collate_samples(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+    return {
+        "images": [item["image"] for item in batch],
+        "targets": torch.stack([item["target"] for item in batch], dim=0),
+        "image_ids": [item["image_id"] for item in batch],
+        "labels": [item["labels"] for item in batch],
+    }
+
+
+def build_dataloader(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle: bool,
+    num_workers: int,
+    pin_memory: bool,
+) -> DataLoader:
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        collate_fn=collate_samples,
+    )
+
+
+def class_pos_weights(label_frequencies: Dict[str, int], label_vocab: Dict[str, int]) -> torch.Tensor:
+    counts = torch.ones(len(label_vocab), dtype=torch.float32)
+    for label, idx in label_vocab.items():
+        counts[idx] = float(max(1, label_frequencies.get(label, 1)))
+    weights = counts.max() / counts
+    return weights
+
+
+def split_paths(processed_root: Path, split_dir_name: str = "splits") -> Dict[str, Path]:
+    split_dir = processed_root / split_dir_name
+    return {
+        "train": split_dir / "train.json",
+        "val": split_dir / "val.json",
+        "test": split_dir / "test.json",
+    }

src/train/encoders.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import timm
+import torch
+from PIL import Image
+from timm.data import create_transform, resolve_data_config
+
+
+class ImageBackbone(torch.nn.Module):
+    def __init__(self, model_name: str, pretrained: bool = True) -> None:
+        super().__init__()
+        self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+        cfg = resolve_data_config(self.model.pretrained_cfg, model=self.model)
+        self.transform = create_transform(**cfg, is_training=False)
+        self.output_dim = getattr(self.model, "num_features", None)
+        if self.output_dim is None:
+            raise RuntimeError(f"Could not infer output dim for {model_name}")
+
+    def forward_pil(self, images: List[Image.Image], device: torch.device) -> torch.Tensor:
+        batch = torch.stack([self.transform(img) for img in images], dim=0).to(device)
+        return self.model(batch)
+
+
+def build_backbones(
+    use_clip: bool,
+    use_dino: bool,
+    clip_model_name: str,
+    dino_model_name: str,
+    clip_pretrained: bool,
+    dino_pretrained: bool,
+) -> Tuple[ImageBackbone | None, ImageBackbone | None, int]:
+    clip = None
+    dino = None
+    feature_dim = 0
+    if use_clip:
+        clip = ImageBackbone(model_name=clip_model_name, pretrained=clip_pretrained)
+        feature_dim += int(clip.output_dim)
+    if use_dino:
+        dino = ImageBackbone(model_name=dino_model_name, pretrained=dino_pretrained)
+        feature_dim += int(dino.output_dim)
+    if feature_dim <= 0:
+        raise RuntimeError("At least one encoder must be enabled for baseline training.")
+    return clip, dino, feature_dim

src/train/losses.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import torch
+
+
+def bce_logits_loss(logits: torch.Tensor, targets: torch.Tensor, pos_weight: torch.Tensor) -> torch.Tensor:
+    loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
+    return loss_fn(logits, targets)

src/train/metrics.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import numpy as np
+from sklearn.metrics import average_precision_score, f1_score
+
+
+def _safe_average_precision(y_true: np.ndarray, y_score: np.ndarray) -> float:
+    if np.unique(y_true).size < 2:
+        return 0.0
+    return float(average_precision_score(y_true, y_score))
+
+
+def compute_metrics(
+    targets: np.ndarray,
+    logits: np.ndarray,
+    threshold: float = 0.5,
+) -> Dict[str, Any]:
+    probs = 1.0 / (1.0 + np.exp(-logits))
+    preds = (probs >= threshold).astype(np.int32)
+    per_attr_ap = []
+    per_attr_f1 = []
+    for idx in range(targets.shape[1]):
+        yt = targets[:, idx]
+        yp = preds[:, idx]
+        ys = probs[:, idx]
+        per_attr_ap.append(_safe_average_precision(yt, ys))
+        per_attr_f1.append(float(f1_score(yt, yp, zero_division=0)))
+
+    macro_f1 = float(f1_score(targets, preds, average="macro", zero_division=0))
+    micro_f1 = float(f1_score(targets, preds, average="micro", zero_division=0))
+    map_score = float(np.mean(per_attr_ap)) if per_attr_ap else 0.0
+    return {
+        "map": map_score,
+        "macro_f1": macro_f1,
+        "micro_f1": micro_f1,
+        "per_attribute_ap": per_attr_ap,
+        "per_attribute_f1": per_attr_f1,
+    }

src/train/models.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import List
+
+import torch
+from PIL import Image
+
+from .encoders import ImageBackbone
+
+
+class BaselineClassifier(torch.nn.Module):
+    def __init__(
+        self,
+        clip_backbone: ImageBackbone | None,
+        dino_backbone: ImageBackbone | None,
+        feature_dim: int,
+        hidden_dims: List[int],
+        dropout: float,
+        num_labels: int,
+    ) -> None:
+        super().__init__()
+        self.clip_backbone = clip_backbone
+        self.dino_backbone = dino_backbone
+        layers: List[torch.nn.Module] = []
+        input_dim = feature_dim
+        for hidden in hidden_dims:
+            layers.append(torch.nn.Linear(input_dim, hidden))
+            layers.append(torch.nn.ReLU())
+            layers.append(torch.nn.Dropout(dropout))
+            input_dim = hidden
+        layers.append(torch.nn.Linear(input_dim, num_labels))
+        self.classifier = torch.nn.Sequential(*layers)
+
+    def encode(self, images: List[Image.Image], device: torch.device) -> torch.Tensor:
+        features = []
+        if self.clip_backbone is not None:
+            features.append(self.clip_backbone.forward_pil(images, device=device))
+        if self.dino_backbone is not None:
+            features.append(self.dino_backbone.forward_pil(images, device=device))
+        if len(features) == 1:
+            return features[0]
+        return torch.cat(features, dim=1)
+
+    def forward(self, images: List[Image.Image], device: torch.device) -> torch.Tensor:
+        feats = self.encode(images=images, device=device)
+        return self.classifier(feats)
+
+    def freeze_backbones(self) -> None:
+        for backbone in (self.clip_backbone, self.dino_backbone):
+            if backbone is None:
+                continue
+            for param in backbone.parameters():
+                param.requires_grad = False