feat(fc2b): add Camie tagger ONNX wrapper
CPU-only, lazy-loaded, process-singleton ONNX session. Parses Camie's string-category selected_tags.csv (vs WD14's integer Danbooru ids). STORE_FLOOR (0.05) keeps the stored predictions JSON compact; SURFACED_CATEGORIES gates which categories the suggestion UI shows (meta/rating/year stored but never surfaced). Inference itself isn't unit-tested (1GB model not in CI); the missing- model error path and pure-logic surface are. Full inference runs in the local integration suite. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
|||||||
|
"""ML pipeline services: tagger, embedder, suggestions, centroids, allowlist, aliases."""
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user