diff --git a/backend/app/services/ml/__init__.py b/backend/app/services/ml/__init__.py new file mode 100644 index 0000000..82ab712 --- /dev/null +++ b/backend/app/services/ml/__init__.py @@ -0,0 +1 @@ +"""ML pipeline services: tagger, embedder, suggestions, centroids, allowlist, aliases.""" diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py new file mode 100644 index 0000000..0ee33eb --- /dev/null +++ b/backend/app/services/ml/tagger.py @@ -0,0 +1,137 @@ +"""Camie-tagger-v2 ONNX wrapper. + +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). +""" + +import csv +import os +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import onnxruntime as ort +from PIL import Image, ImageFile + +# Tolerate minutely-truncated source images (same rationale as IR's wd14.py: +# a few missing bytes at the JPEG EOI shouldn't block tagging the whole image). +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" + +# Below this confidence, predictions aren't stored (keeps the JSON compact). +STORE_FLOOR = float(os.environ.get("TAGGER_STORE_FLOOR", "0.05")) + +# The categories FC-2b surfaces in the UI. Others (meta/rating/year) are +# still stored but the suggestion service filters them out. +SURFACED_CATEGORIES = {"artist", "character", "copyright", "general"} + + +@dataclass(frozen=True) +class TagPrediction: + name: str + category: str + confidence: float + + +class Tagger: + def __init__(self, model_dir: Path | None = None): + self._model_dir = model_dir or _MODEL_DIR + self._session: ort.InferenceSession | None = None + self._tag_meta: list[dict] | None = None + self._input_name: str | None = None + self._output_name: str | None = None + self._input_size: int = 448 + + 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" + if not model_path.is_file(): + raise RuntimeError( + f"Camie model.onnx missing at {model_path}. " + f"Populate /models via the ml-worker downloader." + ) + if not tags_path.is_file(): + raise RuntimeError( + f"Camie selected_tags.csv missing at {tags_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"]} + ) + + session = ort.InferenceSession( + 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._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") + + w, h = img.size + side = max(w, h) + square = Image.new("RGB", (side, side), (255, 255, 255)) + 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 + + 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).""" + self.load() + x = self._preprocess(image_path) + out = self._session.run([self._output_name], {self._input_name: x})[0][0] + results: dict[str, TagPrediction] = {} + for idx, score in enumerate(out): + 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 + ) + return results + + +_default_tagger: Tagger | None = None + + +def get_tagger() -> Tagger: + """Process-level singleton so the ONNX session loads once per worker.""" + global _default_tagger + if _default_tagger is None: + _default_tagger = Tagger() + return _default_tagger diff --git a/tests/test_ml_tagger.py b/tests/test_ml_tagger.py new file mode 100644 index 0000000..de288d7 --- /dev/null +++ b/tests/test_ml_tagger.py @@ -0,0 +1,41 @@ +"""Tagger unit tests. The ONNX model isn't available in CI (it's a 1GB +download into /models), so these test the pure-logic surface: STORE_FLOOR +constant, SURFACED_CATEGORIES set, TagPrediction dataclass, and the +load()-missing-file error path. Full inference is exercised by the local +integration suite against a real /models volume. +""" + +import pytest + +from backend.app.services.ml.tagger import ( + STORE_FLOOR, + SURFACED_CATEGORIES, + TagPrediction, + Tagger, + get_tagger, +) + + +def test_surfaced_categories(): + assert SURFACED_CATEGORIES == {"artist", "character", "copyright", "general"} + + +def test_store_floor_is_low(): + assert 0 < STORE_FLOOR < 0.2 + + +def test_tag_prediction_dataclass(): + p = TagPrediction(name="x", category="general", confidence=0.9) + assert p.name == "x" + assert p.category == "general" + assert p.confidence == 0.9 + + +def test_get_tagger_singleton(): + assert get_tagger() is get_tagger() + + +def test_load_raises_when_model_missing(tmp_path): + t = Tagger(model_dir=tmp_path / "nonexistent") + with pytest.raises(RuntimeError, match="model.onnx missing"): + t.load()