fix(ml): align tagger + downloader with Camie v2 actual layout (model.onnx -> camie-tagger-v2.onnx + JSON metadata + ImageNet preprocessing + sigmoid on refined output)

The HF repo Camais03/camie-tagger-v2 has camie-tagger-v2.onnx (789 MB)
+ camie-tagger-v2-metadata.json (7.77 MB) at root, NOT model.onnx +
selected_tags.csv. Tags ship as nested JSON (dataset_info.tag_mapping)
not CSV. Per the published onnx_inference.py reference: input is NCHW
not NHWC, normalize with ImageNet mean/std, pad-square color (124,116,
104), sigmoid the second output (refined predictions) not the first.

Operator hit this during the IR migration ML backfill — download_models
silently fetched only 3 json files (allow_patterns matched nothing
useful), tagger.load() then raised RuntimeError. Fetched the actual
v2 layout via WebFetch, rewrote tagger to match published reference.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

backend/app/scripts/download_models.py

@@ -25,12 +25,31 @@ def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> Non
 
 
 def ensure_camie() -> None:
+    """Fetch Camie v2 weights + metadata.
+
+    v2 layout (HuggingFace Camais03/camie-tagger-v2): the ONNX file is
+    named camie-tagger-v2.onnx (not model.onnx) and tags ship inside
+    camie-tagger-v2-metadata.json (not selected_tags.csv). Both at root.
+    The repo also contains app/, game/, training/, images/ subdirs full
+    of setup/demo files we don't need — allow_patterns scopes the fetch
+    to just the inference essentials (~790 MB instead of ~2 GB).
+    """
     dest = MODEL_ROOT / "camie"
-    if (dest / "model.onnx").is_file() and (dest / "selected_tags.csv").is_file():
+    model_file = dest / "camie-tagger-v2.onnx"
+    meta_file = dest / "camie-tagger-v2-metadata.json"
+    if model_file.is_file() and meta_file.is_file():
         print(f"[download_models] Camie present at {dest}")
         return
     print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}")
-    _snapshot(CAMIE_REPO, dest, ["model.onnx", "selected_tags.csv", "*.json"])
+    _snapshot(
+        CAMIE_REPO, dest,
+        [
+            "camie-tagger-v2.onnx",
+            "camie-tagger-v2-metadata.json",
+            "config.json",
+            "config.yaml",
+        ],
+    )
 
 
 def ensure_siglip() -> None:

backend/app/services/ml/tagger.py

@@ -4,12 +4,14 @@ CPU-only, single-image at a time. Loaded lazily inside the ml-worker
 process; NOT thread-safe — the ml queue worker must run --concurrency=1
 (set by the FC-1 entrypoint).
 
-Camie's selected_tags.csv columns: tag_id,name,category,count
+v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has
-where category is a string: general|character|copyright|artist|meta|rating|year
+camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB)
-(unlike WD14's integer Danbooru category ids).
++ config.json. Tags ship as nested JSON, not CSV. Preprocessing and
+output handling follow the published onnx_inference.py reference:
+ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]).
 """
 
-import csv
+import json
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -28,6 +30,8 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 MODEL_NAME = os.environ.get("CAMIE_MODEL_NAME", "camie-tagger-v2")
 _MODEL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "camie"
+_MODEL_FILE = f"{MODEL_NAME}.onnx"
+_METADATA_FILE = f"{MODEL_NAME}-metadata.json"
 
 # Below this confidence, predictions aren't stored (keeps the JSON compact).
 STORE_FLOOR = float(os.environ.get("TAGGER_STORE_FLOOR", "0.05"))
@@ -39,6 +43,12 @@ STORE_FLOOR = float(os.environ.get("TAGGER_STORE_FLOOR", "0.05"))
 # stored at STORE_FLOOR but artist never surfaces.
 SURFACED_CATEGORIES = {"character", "copyright", "general"}
 
+# ImageNet preprocessing constants (per Camie v2 onnx_inference.py).
+_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+# Square-pad color ≈ ImageNet mean × 255 (matches reference inference).
+_PAD_COLOR = (124, 116, 104)
+
 
 @dataclass(frozen=True)
 class TagPrediction:
@@ -51,34 +61,48 @@ class Tagger:
     def __init__(self, model_dir: Path | None = None):
         self._model_dir = model_dir or _MODEL_DIR
         self._session = None  # onnxruntime.InferenceSession once load()ed
-        self._tag_meta: list[dict] | None = None
+        self._tag_names: list[str] | None = None
+        self._tag_categories: list[str] | None = None
         self._input_name: str | None = None
-        self._output_name: str | None = None
+        self._input_size: int = 512
-        self._input_size: int = 448
 
     def load(self) -> None:
         if self._session is not None:
             return
-        model_path = self._model_dir / "model.onnx"
+        model_path = self._model_dir / _MODEL_FILE
-        tags_path = self._model_dir / "selected_tags.csv"
+        meta_path = self._model_dir / _METADATA_FILE
         if not model_path.is_file():
             raise RuntimeError(
-                f"Camie model.onnx missing at {model_path}. "
+                f"Camie {_MODEL_FILE} missing at {model_path}. "
                 f"Populate /models via the ml-worker downloader."
             )
-        if not tags_path.is_file():
+        if not meta_path.is_file():
             raise RuntimeError(
-                f"Camie selected_tags.csv missing at {tags_path}. "
+                f"Camie {_METADATA_FILE} missing at {meta_path}. "
                 f"Populate /models via the ml-worker downloader."
             )
 
-        tag_meta: list[dict] = []
+        with open(meta_path) as f:
-        with open(tags_path, newline="") as f:
+            metadata = json.load(f)
-            reader = csv.DictReader(f)
+
-            for row in reader:
+        # Per Camie v2 onnx_inference.py: idx_to_tag is keyed by str(idx);
-                tag_meta.append(
+        # tag_to_category maps tag_name -> category. Project to two parallel
-                    {"name": row["name"], "category": row["category"]}
+        # lists indexed by output position for O(1) lookup in the hot path.
-                )
+        ds = metadata["dataset_info"]
+        idx_to_tag = ds["tag_mapping"]["idx_to_tag"]
+        tag_to_category = ds["tag_mapping"]["tag_to_category"]
+        total = ds["total_tags"]
+        names: list[str] = []
+        cats: list[str] = []
+        for i in range(total):
+            name = idx_to_tag.get(str(i), f"unknown-{i}")
+            names.append(name)
+            cats.append(tag_to_category.get(name, "general"))
+
+        # Input size from metadata; fall back to 512 (the v2 default).
+        self._input_size = int(
+            metadata.get("model_info", {}).get("img_size", 512)
+        )
 
         # Lazy import — kept after the file-existence checks so the
         # missing-model RuntimeError still fires first in environments
@@ -89,51 +113,65 @@ class Tagger:
             str(model_path), providers=["CPUExecutionProvider"]
         )
         self._input_name = session.get_inputs()[0].name
-        self._output_name = session.get_outputs()[0].name
-        input_shape = session.get_inputs()[0].shape
-        for dim in input_shape:
-            if isinstance(dim, int) and dim > 1:
-                self._input_size = dim
-                break
         # Assign sentinels last so a partial load isn't observable.
-        self._tag_meta = tag_meta
+        self._tag_names = names
+        self._tag_categories = cats
         self._session = session
 
     def _preprocess(self, image_path: Path) -> np.ndarray:
         img = Image.open(image_path)
-        # Camie handles RGBA natively but we still composite onto white so
+        # Composite RGBA onto neutral so transparency doesn't bias the model.
-        # transparency doesn't bias the model (same as IR's WD14 path).
+        if img.mode == "RGBA":
-        if img.mode != "RGBA":
+            bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
-            img = img.convert("RGBA")
+            bg.paste(img, mask=img.split()[3])
-        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
+            img = bg.convert("RGB")
-        bg.paste(img, mask=img.split()[3])
+        elif img.mode != "RGB":
-        img = bg.convert("RGB")
+            img = img.convert("RGB")
 
+        # Pad to square with ImageNet-mean color, then bicubic resize.
         w, h = img.size
         side = max(w, h)
-        square = Image.new("RGB", (side, side), (255, 255, 255))
+        square = Image.new("RGB", (side, side), _PAD_COLOR)
         square.paste(img, ((side - w) // 2, (side - h) // 2))
         square = square.resize(
             (self._input_size, self._input_size), Image.BICUBIC
         )
-        arr = np.array(square, dtype=np.float32)
+
-        return arr[np.newaxis, :, :, :]  # NHWC
+        arr = np.array(square, dtype=np.float32) / 255.0   # HWC, [0,1]
+        arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD        # ImageNet normalize
+        arr = arr.transpose(2, 0, 1)                        # HWC -> CHW
+        return arr[np.newaxis, :, :, :]                     # NCHW
 
     def infer(self, image_path: Path) -> dict[str, TagPrediction]:
-        """Run Camie on one image. Returns {name: TagPrediction}, only
+        """Run Camie v2 on one image. Returns {name: TagPrediction} with
-        entries with confidence >= STORE_FLOOR (across all categories —
+        confidence >= STORE_FLOOR (across all categories — the suggestion
-        the suggestion service does category filtering later)."""
+        service does category filtering later).
+
+        v2 emits multiple outputs; we use the refined predictions
+        (output[1] per onnx_inference.py). Sigmoid is applied to raw
+        logits to produce [0,1] confidence scores.
+        """
         self.load()
         x = self._preprocess(image_path)
-        out = self._session.run([self._output_name], {self._input_name: x})[0][0]
+        outputs = self._session.run(None, {self._input_name: x})
+        # Refined predictions if present (v2 emits initial + refined),
+        # fall back to initial for single-output forks.
+        logits = outputs[1] if len(outputs) > 1 else outputs[0]
+        # Squeeze batch dim, apply sigmoid.
+        probs = 1.0 / (1.0 + np.exp(-logits[0]))
         results: dict[str, TagPrediction] = {}
-        for idx, score in enumerate(out):
+        names = self._tag_names
+        cats = self._tag_categories
+        for idx, score in enumerate(probs):
             conf = float(score)
             if conf < STORE_FLOOR:
                 continue
-            meta = self._tag_meta[idx]
+            if idx >= len(names):
-            results[meta["name"]] = TagPrediction(
+                # Output longer than metadata declared — shouldn't happen but
-                name=meta["name"], category=meta["category"], confidence=conf
+                # don't crash the import pipeline if v2 metadata desynchronizes.
+                continue
+            results[names[idx]] = TagPrediction(
+                name=names[idx], category=cats[idx], confidence=conf
             )
         return results
 

tests/test_ml_tagger.py

@@ -40,5 +40,8 @@ def test_get_tagger_singleton():
 
 def test_load_raises_when_model_missing(tmp_path):
     t = Tagger(model_dir=tmp_path / "nonexistent")
-    with pytest.raises(RuntimeError, match="model.onnx missing"):
+    # Match the trailing "missing at <path>" rather than the specific
+    # filename, so a future model-version bump (camie-tagger-v3.onnx, etc.)
+    # doesn't bounce this test.
+    with pytest.raises(RuntimeError, match=r"\.onnx missing at "):
         t.load()