From 3b3e7565fb65dcc351c2a53040ad91c432caa275 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Mon, 25 May 2026 02:25:30 -0400 Subject: [PATCH] fix(ml): align tagger + downloader with Camie v2 actual layout (model.onnx -> camie-tagger-v2.onnx + JSON metadata + ImageNet preprocessing + sigmoid on refined output) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- backend/app/scripts/download_models.py | 23 ++++- backend/app/services/ml/tagger.py | 126 ++++++++++++++++--------- tests/test_ml_tagger.py | 5 +- 3 files changed, 107 insertions(+), 47 deletions(-) diff --git a/backend/app/scripts/download_models.py b/backend/app/scripts/download_models.py index 39131a1..4b5e6d3 100644 --- a/backend/app/scripts/download_models.py +++ b/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: diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py index 92de3a3..5277b44 100644 --- a/backend/app/services/ml/tagger.py +++ b/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 -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 diff --git a/tests/test_ml_tagger.py b/tests/test_ml_tagger.py index 79cfe51..17c449c 100644 --- a/tests/test_ml_tagger.py +++ b/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 " 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()