feat(agent): crop proposers — booru_yolo anatomy + COCO person + comic panels (#1202)
Better region PROPOSERS feeding the existing crop→SigLIP→max-over-bag heads (no change to the learned-tagging approach; no per-tag cost — propose once, embed each region, all heads in one matmul). - detectors.py: lazy ultralytics YOLO wrapper, each proposer independently optional + guarded (a bad weight spec / inference error self-disables that one, logged, never breaks the worker). Weights resolve from an ultralytics name | http(s) URL | "hf_repo::file", cached under HF_HOME. NMS merge so a figure two detectors both find collapses to one crop. - worker: figure boxes = imgutils detect_person ∪ general COCO person (merged) → CCIP + concept (anime + Western/realistic coverage); booru_yolo anatomy components (head/cat-head/anatomy/…) → concept crops; comic panels → kind= 'panel' concept crops. Capped per frame (MAX_COMPONENTS/MAX_PANELS). - config + compose: PERSON_WEIGHTS (default yolo11n.pt, works OOB), ANATOMY_WEIGHTS + PANEL_WEIGHTS (operator sets booru_yolo URL + mosesb panel hf::file; empty = off). ultralytics added to requirements. - backend: image_region 'kind' doc notes 'panel'; no migration (free String, and the bag scorer keys on a non-null siglip_embedding, not the kind, so any SigLIP region joins the bag automatically). Agent is outside CI — py-compiled here; operator tests on the GPU and checks Western-vs-anime crop quality via /api/ccip observability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
@@ -37,6 +37,17 @@ services:
|
|||||||
# Crop embedder (SigLIP concept bag): float16 keeps VRAM low on a shared
|
# Crop embedder (SigLIP concept bag): float16 keeps VRAM low on a shared
|
||||||
# desktop GPU; the model itself is announced by the server.
|
# desktop GPU; the model itself is announced by the server.
|
||||||
SIGLIP_DTYPE: ${SIGLIP_DTYPE:-float16}
|
SIGLIP_DTYPE: ${SIGLIP_DTYPE:-float16}
|
||||||
|
# Crop PROPOSERS (extra YOLO detectors → more/better concept crops).
|
||||||
|
# PERSON_WEIGHTS: general COCO person detector (Western/realistic figures),
|
||||||
|
# merged with the anime person detector. Default works out of the box.
|
||||||
|
# ANATOMY_WEIGHTS: booru_yolo (anime/furry/NSFW components). Set to the
|
||||||
|
# weights URL, e.g. the yolov11m_aa22.pt from github.com/aperveyev/booru_yolo.
|
||||||
|
# PANEL_WEIGHTS: comic-panel detector as "hf_repo::file",
|
||||||
|
# e.g. mosesb/best-comic-panel-detection::<weights>.pt
|
||||||
|
# Empty = that proposer is off. Each self-disables if its weights fail to load.
|
||||||
|
PERSON_WEIGHTS: ${PERSON_WEIGHTS:-yolo11n.pt}
|
||||||
|
ANATOMY_WEIGHTS: ${ANATOMY_WEIGHTS:-}
|
||||||
|
PANEL_WEIGHTS: ${PANEL_WEIGHTS:-}
|
||||||
volumes:
|
volumes:
|
||||||
# Persist the downloaded ONNX models so restarts are fast.
|
# Persist the downloaded ONNX models so restarts are fast.
|
||||||
- fc-agent-models:/models
|
- fc-agent-models:/models
|
||||||
|
|||||||
@@ -18,6 +18,16 @@ class Config:
|
|||||||
# the server announces in the lease)
|
# the server announces in the lease)
|
||||||
auto_start: bool # start the worker pool on boot (so a container restart
|
auto_start: bool # start the worker pool on boot (so a container restart
|
||||||
# resumes processing without anyone clicking Start)
|
# resumes processing without anyone clicking Start)
|
||||||
|
# Crop PROPOSERS (extra YOLO detectors that say where to crop). Each weight
|
||||||
|
# spec is an ultralytics name | http(s) URL | "hf_repo::file" ("" = off).
|
||||||
|
person_weights: str # general COCO person detector (Western/realistic figs)
|
||||||
|
person_conf: float
|
||||||
|
anatomy_weights: str # booru_yolo anime/furry/NSFW components
|
||||||
|
anatomy_conf: float
|
||||||
|
panel_weights: str # comic-panel detector
|
||||||
|
panel_conf: float
|
||||||
|
max_components: int # cap anatomy component crops per frame
|
||||||
|
max_panels: int # cap panel crops per frame
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_env(cls) -> "Config":
|
def from_env(cls) -> "Config":
|
||||||
@@ -33,4 +43,12 @@ class Config:
|
|||||||
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
||||||
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
||||||
auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"),
|
auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"),
|
||||||
|
person_weights=os.environ.get("PERSON_WEIGHTS", "yolo11n.pt"),
|
||||||
|
person_conf=float(os.environ.get("PERSON_CONF", "0.35")),
|
||||||
|
anatomy_weights=os.environ.get("ANATOMY_WEIGHTS", ""),
|
||||||
|
anatomy_conf=float(os.environ.get("ANATOMY_CONF", "0.30")),
|
||||||
|
panel_weights=os.environ.get("PANEL_WEIGHTS", ""),
|
||||||
|
panel_conf=float(os.environ.get("PANEL_CONF", "0.30")),
|
||||||
|
max_components=int(os.environ.get("MAX_COMPONENTS", "8")),
|
||||||
|
max_panels=int(os.environ.get("MAX_PANELS", "8")),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -0,0 +1,163 @@
|
|||||||
|
"""Region PROPOSERS — small YOLO detectors that decide WHERE to crop. They run
|
||||||
|
on the agent GPU and their boxes feed the crop → SigLIP → max-over-bag pipeline:
|
||||||
|
|
||||||
|
- person (general COCO yolo11n): full-figure boxes for realistic / Western art
|
||||||
|
the anime person-detector misses; NMS-merged with imgutils detect_person and
|
||||||
|
fed to CCIP (identity) + a concept crop.
|
||||||
|
- anatomy (booru_yolo): anime / furry / NSFW torso components (head, cat-head,
|
||||||
|
boob, hip, …) — concept crops aligned to the operator's tag vocabulary.
|
||||||
|
- panel (mosesb): a comic page → panel regions → concept crops.
|
||||||
|
|
||||||
|
Each proposer is INDEPENDENTLY optional + guarded: a bad weight path or an
|
||||||
|
inference error disables just that proposer (logged) and never breaks the
|
||||||
|
worker, which still falls back to imgutils detection. Weights resolve from an
|
||||||
|
ultralytics builtin name ("yolo11n.pt"), an http(s) URL, or "hf_repo::file" —
|
||||||
|
cached under HF_HOME so the download happens once.
|
||||||
|
"""
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import threading
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
log = logging.getLogger("fc_agent.detectors")
|
||||||
|
_CACHE = Path(os.environ.get("HF_HOME", "/models")) / "yolo"
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve(spec: str) -> str | None:
|
||||||
|
"""A local weights path (downloading if needed) or an ultralytics builtin
|
||||||
|
name. None if the spec is empty/unresolvable."""
|
||||||
|
if not spec:
|
||||||
|
return None
|
||||||
|
if "::" in spec: # hf_repo::filename
|
||||||
|
repo, _, fname = spec.partition("::")
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
return hf_hub_download(
|
||||||
|
repo_id=repo, filename=fname, cache_dir=str(_CACHE)
|
||||||
|
)
|
||||||
|
if spec.startswith(("http://", "https://")):
|
||||||
|
_CACHE.mkdir(parents=True, exist_ok=True)
|
||||||
|
dest = _CACHE / spec.rsplit("/", 1)[-1]
|
||||||
|
if not dest.is_file():
|
||||||
|
import requests
|
||||||
|
r = requests.get(spec, timeout=300)
|
||||||
|
r.raise_for_status()
|
||||||
|
dest.write_bytes(r.content)
|
||||||
|
return str(dest)
|
||||||
|
return spec # ultralytics builtin name
|
||||||
|
|
||||||
|
|
||||||
|
def _iou(a, b) -> float:
|
||||||
|
ax, ay, aw, ah = a
|
||||||
|
bx, by, bw, bh = b
|
||||||
|
ix = max(0.0, min(ax + aw, bx + bw) - max(ax, bx))
|
||||||
|
iy = max(0.0, min(ay + ah, by + bh) - max(ay, by))
|
||||||
|
inter = ix * iy
|
||||||
|
union = aw * ah + bw * bh - inter
|
||||||
|
return inter / union if union > 0 else 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def nms_merge(boxes, iou_thresh: float = 0.6):
|
||||||
|
"""Greedy NMS over (bbox_norm, score, label) from possibly several detectors,
|
||||||
|
so the same figure found by two of them collapses to one (higher-score) box."""
|
||||||
|
kept = []
|
||||||
|
for bb, sc, lb in sorted(boxes, key=lambda b: b[1], reverse=True):
|
||||||
|
if all(_iou(bb, k[0]) < iou_thresh for k in kept):
|
||||||
|
kept.append((bb, sc, lb))
|
||||||
|
return kept
|
||||||
|
|
||||||
|
|
||||||
|
class YoloProposer:
|
||||||
|
"""One lazily-loaded ultralytics YOLO. detect(image) → [(bbox_norm, score,
|
||||||
|
label)] with bbox normalized (x, y, w, h) in [0,1]. Self-disables on any
|
||||||
|
load/inference failure."""
|
||||||
|
|
||||||
|
def __init__(self, name, weights, conf=0.25, keep_labels=None):
|
||||||
|
self.name = name
|
||||||
|
self._spec = weights
|
||||||
|
self._conf = conf
|
||||||
|
self._keep = [k.lower() for k in keep_labels] if keep_labels else None
|
||||||
|
self._model = None
|
||||||
|
self._ok = True
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
|
def _load(self):
|
||||||
|
if self._model is not None or not self._ok:
|
||||||
|
return
|
||||||
|
with self._lock:
|
||||||
|
if self._model is not None or not self._ok:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
from ultralytics import YOLO
|
||||||
|
path = _resolve(self._spec)
|
||||||
|
if path is None:
|
||||||
|
self._ok = False
|
||||||
|
return
|
||||||
|
self._model = YOLO(path)
|
||||||
|
log.info("detector %s loaded (%s)", self.name, path)
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
log.warning("detector %s disabled (load failed): %s", self.name, exc)
|
||||||
|
self._ok = False
|
||||||
|
|
||||||
|
def detect(self, image):
|
||||||
|
self._load()
|
||||||
|
if self._model is None:
|
||||||
|
return []
|
||||||
|
try:
|
||||||
|
res = self._model.predict(image, conf=self._conf, verbose=False)[0]
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
log.warning("detector %s inference failed: %s", self.name, exc)
|
||||||
|
return []
|
||||||
|
iw, ih = image.size
|
||||||
|
names = getattr(res, "names", None) or {}
|
||||||
|
out = []
|
||||||
|
for b in res.boxes:
|
||||||
|
label = str(names.get(int(b.cls), int(b.cls))).lower()
|
||||||
|
if self._keep is not None and not any(k in label for k in self._keep):
|
||||||
|
continue
|
||||||
|
x0, y0, x1, y1 = (float(v) for v in b.xyxy[0].tolist())
|
||||||
|
out.append((
|
||||||
|
(x0 / iw, y0 / ih, (x1 - x0) / iw, (y1 - y0) / ih),
|
||||||
|
float(b.conf), label,
|
||||||
|
))
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class Proposers:
|
||||||
|
"""The agent's proposer set, built from config. Each detector is optional —
|
||||||
|
an empty weight spec leaves that proposer off."""
|
||||||
|
|
||||||
|
def __init__(self, cfg):
|
||||||
|
self.cfg = cfg
|
||||||
|
self._person = (
|
||||||
|
YoloProposer("person-coco", cfg.person_weights,
|
||||||
|
conf=cfg.person_conf, keep_labels=["person"])
|
||||||
|
if cfg.person_weights else None
|
||||||
|
)
|
||||||
|
self._anatomy = (
|
||||||
|
YoloProposer("anatomy", cfg.anatomy_weights, conf=cfg.anatomy_conf)
|
||||||
|
if cfg.anatomy_weights else None
|
||||||
|
)
|
||||||
|
self._panel = (
|
||||||
|
YoloProposer("panel", cfg.panel_weights, conf=cfg.panel_conf)
|
||||||
|
if cfg.panel_weights else None
|
||||||
|
)
|
||||||
|
|
||||||
|
def figures(self, image, base_boxes):
|
||||||
|
"""Merge imgutils person boxes (base_boxes: [(bbox, score)]) with the
|
||||||
|
general COCO person detector → NMS'd figure boxes [(bbox, score, label)]."""
|
||||||
|
boxes = [(bb, sc if sc is not None else 1.0, "person") for bb, sc in base_boxes]
|
||||||
|
if self._person is not None:
|
||||||
|
boxes += self._person.detect(image)
|
||||||
|
return nms_merge(boxes)
|
||||||
|
|
||||||
|
def components(self, image):
|
||||||
|
if self._anatomy is None:
|
||||||
|
return []
|
||||||
|
items = sorted(self._anatomy.detect(image), key=lambda b: b[1], reverse=True)
|
||||||
|
return items[: self.cfg.max_components]
|
||||||
|
|
||||||
|
def panels(self, image):
|
||||||
|
if self._panel is None:
|
||||||
|
return []
|
||||||
|
items = sorted(self._panel.detect(image), key=lambda b: b[1], reverse=True)
|
||||||
|
return items[: self.cfg.max_panels]
|
||||||
+51
-14
@@ -76,6 +76,9 @@ class Worker:
|
|||||||
# needs it, from the model the server announces — one shared instance.
|
# needs it, from the model the server announces — one shared instance.
|
||||||
self._embedder = None
|
self._embedder = None
|
||||||
self._embedder_lock = threading.Lock()
|
self._embedder_lock = threading.Lock()
|
||||||
|
# Region proposers (extra YOLO detectors) — lazily built once, shared.
|
||||||
|
self._proposers = None
|
||||||
|
self._proposers_lock = threading.Lock()
|
||||||
|
|
||||||
# --- control -----------------------------------------------------------
|
# --- control -----------------------------------------------------------
|
||||||
def start(self):
|
def start(self):
|
||||||
@@ -177,6 +180,15 @@ class Worker:
|
|||||||
self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
|
self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
|
||||||
return self._embedder
|
return self._embedder
|
||||||
|
|
||||||
|
def _ensure_proposers(self):
|
||||||
|
if self._proposers is not None:
|
||||||
|
return self._proposers
|
||||||
|
with self._proposers_lock:
|
||||||
|
if self._proposers is None:
|
||||||
|
from .detectors import Proposers
|
||||||
|
self._proposers = Proposers(self.cfg)
|
||||||
|
return self._proposers
|
||||||
|
|
||||||
def _process(self, job: dict) -> bool:
|
def _process(self, job: dict) -> bool:
|
||||||
"""Process one job. Returns True when handled (completed, or hard-failed
|
"""Process one job. Returns True when handled (completed, or hard-failed
|
||||||
because the job itself is bad) and False on a TRANSPORT error (curator
|
because the job itself is bad) and False on a TRANSPORT error (curator
|
||||||
@@ -227,43 +239,68 @@ class Worker:
|
|||||||
want_ccip = task in ("ccip", "both")
|
want_ccip = task in ("ccip", "both")
|
||||||
want_siglip = task in ("ccip", "siglip", "both")
|
want_siglip = task in ("ccip", "siglip", "both")
|
||||||
replace_kinds = (
|
replace_kinds = (
|
||||||
["concept"] if task == "siglip" else ["figure", "face", "concept"]
|
["concept", "panel"] if task == "siglip"
|
||||||
|
else ["figure", "face", "concept", "panel"]
|
||||||
)
|
)
|
||||||
embedder = self._ensure_embedder(model_name) if want_siglip else None
|
embedder = self._ensure_embedder(model_name) if want_siglip else None
|
||||||
|
proposers = self._ensure_proposers()
|
||||||
|
|
||||||
regions = []
|
regions = []
|
||||||
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
||||||
dv = f"person-{self.cfg.detector_level}"
|
dv = f"person-{self.cfg.detector_level}"
|
||||||
|
|
||||||
|
def _concept(frame, bbox, t, score, detver, kind="concept"):
|
||||||
|
"""A SigLIP region for one crop (None if below the size floor)."""
|
||||||
|
crop = crop_region(frame, bbox)
|
||||||
|
if crop is None:
|
||||||
|
return None
|
||||||
|
return {
|
||||||
|
"kind": kind, "bbox": list(bbox), "frame_time": t,
|
||||||
|
"score": score, "siglip_embedding": embedder.embed(crop),
|
||||||
|
"embedding_version": embed_version, "detector_version": detver,
|
||||||
|
}
|
||||||
|
|
||||||
for t, frame in frames:
|
for t, frame in frames:
|
||||||
figs = models.detect_figures(frame, self.cfg.detector_level)
|
# FIGURE boxes: imgutils detect_person ∪ general COCO person,
|
||||||
|
# NMS-merged → CCIP identity (+ a concept crop). Covers anime +
|
||||||
|
# Western/realistic figures.
|
||||||
|
base = models.detect_figures(frame, self.cfg.detector_level)
|
||||||
|
figs = proposers.figures(frame, base)
|
||||||
if not figs:
|
if not figs:
|
||||||
figs = [((0.0, 0.0, 1.0, 1.0), None)] # whole-frame fallback
|
figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")] # whole-frame fallback
|
||||||
for bbox, score in figs:
|
for bbox, score, _label in figs:
|
||||||
crop = crop_region(frame, bbox)
|
crop = crop_region(frame, bbox)
|
||||||
if crop is None:
|
if crop is None:
|
||||||
continue
|
continue
|
||||||
if want_ccip:
|
if want_ccip:
|
||||||
regions.append({
|
regions.append({
|
||||||
"kind": "figure",
|
"kind": "figure", "bbox": list(bbox), "frame_time": t,
|
||||||
"bbox": list(bbox),
|
|
||||||
"frame_time": t,
|
|
||||||
"score": score,
|
"score": score,
|
||||||
"ccip_embedding": models.ccip_vector(
|
"ccip_embedding": models.ccip_vector(
|
||||||
crop, self.cfg.ccip_model or None
|
crop, self.cfg.ccip_model or None
|
||||||
),
|
),
|
||||||
"embedding_version": ccip_ev,
|
"embedding_version": ccip_ev, "detector_version": dv,
|
||||||
"detector_version": dv,
|
|
||||||
})
|
})
|
||||||
if want_siglip:
|
if want_siglip:
|
||||||
regions.append({
|
regions.append({
|
||||||
"kind": "concept",
|
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||||||
"bbox": list(bbox),
|
|
||||||
"frame_time": t,
|
|
||||||
"score": score,
|
"score": score,
|
||||||
"siglip_embedding": embedder.embed(crop),
|
"siglip_embedding": embedder.embed(crop),
|
||||||
"embedding_version": embed_version,
|
"embedding_version": embed_version, "detector_version": dv,
|
||||||
"detector_version": dv,
|
|
||||||
})
|
})
|
||||||
|
if not want_siglip:
|
||||||
|
continue
|
||||||
|
# ANATOMY components (booru_yolo: head/cat-head/anatomy/…) →
|
||||||
|
# concept crops only (not full characters, so no CCIP).
|
||||||
|
for bbox, score, label in proposers.components(frame):
|
||||||
|
r = _concept(frame, bbox, t, score, f"booru:{label}")
|
||||||
|
if r is not None:
|
||||||
|
regions.append(r)
|
||||||
|
# PANEL crops (comic page → panels) → kind='panel' (still SigLIP).
|
||||||
|
for bbox, score, _label in proposers.panels(frame):
|
||||||
|
r = _concept(frame, bbox, t, score, "panel", kind="panel")
|
||||||
|
if r is not None:
|
||||||
|
regions.append(r)
|
||||||
self.client.submit(job["job_id"], regions, replace_kinds)
|
self.client.submit(job["job_id"], regions, replace_kinds)
|
||||||
self._bump(processed=1)
|
self._bump(processed=1)
|
||||||
return True
|
return True
|
||||||
|
|||||||
@@ -7,6 +7,9 @@ onnxruntime-gpu
|
|||||||
# Dockerfile from the CUDA-12.4 wheel index so the GPU build is deterministic;
|
# Dockerfile from the CUDA-12.4 wheel index so the GPU build is deterministic;
|
||||||
# transformers loads whatever SigLIP-family model the server announces.
|
# transformers loads whatever SigLIP-family model the server announces.
|
||||||
transformers>=4.45
|
transformers>=4.45
|
||||||
|
# Crop PROPOSERS — small YOLO detectors (booru_yolo anatomy, COCO person, comic
|
||||||
|
# panel) that decide where to crop. Uses the torch already installed above.
|
||||||
|
ultralytics>=8.3
|
||||||
# Control surface + HTTP.
|
# Control surface + HTTP.
|
||||||
fastapi
|
fastapi
|
||||||
uvicorn[standard]
|
uvicorn[standard]
|
||||||
|
|||||||
@@ -31,7 +31,10 @@ class ImageRegion(Base):
|
|||||||
ForeignKey("image_record.id", ondelete="CASCADE"), index=True
|
ForeignKey("image_record.id", ondelete="CASCADE"), index=True
|
||||||
)
|
)
|
||||||
# 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP
|
# 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP
|
||||||
# character id) | 'concept' (→ SigLIP head bag).
|
# character id) | 'concept' (→ SigLIP head bag) | 'panel' (a comic panel crop,
|
||||||
|
# also SigLIP → the bag). Free String, not an enum — proposers can add kinds
|
||||||
|
# without a migration; the bag scorer keys on a non-null siglip_embedding, not
|
||||||
|
# the kind, so any SigLIP-embedded region joins the bag.
|
||||||
kind: Mapped[str] = mapped_column(String(16), nullable=False)
|
kind: Mapped[str] = mapped_column(String(16), nullable=False)
|
||||||
# For video/animated media: the source frame's timestamp in SECONDS. NULL for
|
# For video/animated media: the source frame's timestamp in SECONDS. NULL for
|
||||||
# static images. Lets a video be a BAG of per-frame instances (fixes the
|
# static images. Lets a video be a BAG of per-frame instances (fixes the
|
||||||
|
|||||||
Reference in New Issue
Block a user