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>
2026-05-25 02:25:30 -04:00
parent 9d5abb09f6
commit 3b3e7565fb
3 changed files with 107 additions and 47 deletions
@@ -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:
@@ -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
-where category is a string: general|character|copyright|artist|meta|rating|year
-(unlike WD14's integer Danbooru category ids).
+v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has
+camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB)
+ 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 = 448
+        self._input_size: int = 512

    def load(self) -> None:
        if self._session is not None:
            return
-        model_path = self._model_dir / "model.onnx"
-        tags_path = self._model_dir / "selected_tags.csv"
+        model_path = self._model_dir / _MODEL_FILE
+        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(tags_path, newline="") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                tag_meta.append(
-                    {"name": row["name"], "category": row["category"]}
-                )
+        with open(meta_path) as f:
+            metadata = json.load(f)
+
+        # Per Camie v2 onnx_inference.py: idx_to_tag is keyed by str(idx);
+        # tag_to_category maps tag_name -> category. Project to two parallel
+        # 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
-        # transparency doesn't bias the model (same as IR's WD14 path).
-        if img.mode != "RGBA":
-            img = img.convert("RGBA")
-        bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
-        bg.paste(img, mask=img.split()[3])
-        img = bg.convert("RGB")
+        # Composite RGBA onto neutral so transparency doesn't bias the model.
+        if img.mode == "RGBA":
+            bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
+            bg.paste(img, mask=img.split()[3])
+            img = bg.convert("RGB")
+        elif img.mode != "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
-        entries with confidence >= STORE_FLOOR (across all categories —
-        the suggestion service does category filtering later)."""
+        """Run Camie v2 on one image. Returns {name: TagPrediction} with
+        confidence >= STORE_FLOOR (across all categories — the suggestion
+        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]
-            results[meta["name"]] = TagPrediction(
-                name=meta["name"], category=meta["category"], confidence=conf
+            if idx >= len(names):
+                # Output longer than metadata declared — shouldn't happen but
+                # 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

@@ -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()