FabledCurator/agent/fc_agent/detectors.py

"""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
import types
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


def dedupe_crops(pending, iou_thresh: float = 0.85):
    """Greedy high-IoU dedupe over a list of (crop, region_template) pairs, run
    just before the batched SigLIP embed so we never embed the same region twice.

    Figure boxes are already NMS-merged and each YOLO self-NMSes, but the combined
    per-frame pile (figure→concept ∪ anatomy component→concept ∪ panel) can still
    carry genuine near-duplicates across proposers — e.g. a figure box that nearly
    coincides with an anatomy component on a solo bust, or overlapping booru head
    classes on one head. Those embed the same region twice, wasting GPU and a slot
    against max_regions.

    Boxes are compared ONLY within the same output kind and dropped when they
    overlap at >= iou_thresh, keeping the highest-scoring one. The HIGH default
    threshold is deliberate: it collapses only true near-identical boxes while
    preserving intentional nested crops across scopes (a whole figure vs a small
    head component sit well below it) and distinct kinds (concept vs panel). A
    value >= 1.0 effectively disables it (nothing but an exact box matches)."""
    kept = []
    kept_boxes: dict = {}   # kind -> [bbox, ...] already kept
    for crop, tmpl in sorted(
        pending, key=lambda p: p[1].get("score") or 0.0, reverse=True
    ):
        bb = tmpl.get("bbox")
        prior = kept_boxes.setdefault(tmpl.get("kind"), [])
        if bb is not None and any(_iou(bb, kb) >= iou_thresh for kb in prior):
            continue
        prior.append(bb)
        kept.append((crop, tmpl))
    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)
                # Disable ultralytics' load-time Conv+BN fusion. AutoBackend fuses
                # the graph on the first predict; some checkpoints (yolo11n, the
                # comic-panel model) crash that step with "'Conv' object has no
                # attribute 'bn'" (a partially-fused / version-mismatched graph),
                # which silently disabled those proposers (operator-flagged
                # 2026-07-01). Unfused inference is correct — only marginally
                # slower — and this is robust across ultralytics versions; if a
                # future version ignores the override, the detect() guard below
                # still self-disables the proposer instead of spamming per image.
                inner = getattr(self._model, "model", None)
                if inner is not None:
                    inner.fuse = types.MethodType(lambda self, *a, **k: self, inner)
                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
            # Permanently self-disable on the FIRST inference failure rather than
            # re-throwing (and re-logging) on every image forever — an unfixable
            # model fault degrades to "this proposer is off", logged once.
            log.warning("detector %s disabled (inference failed): %s", self.name, exc)
            self._ok = False
            self._model = None
            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

    def unload(self) -> None:
        """Drop the loaded YOLO so its VRAM can be reclaimed; detect() reloads it
        lazily on the next job. Leaves _ok untouched — a healthy proposer comes
        back, but one that self-disabled on a fault stays off."""
        with self._lock:
            self._model = None


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)],
        capped to the highest-scoring max_figures. Uncapped, a busy/huge image
        (many characters) yields hundreds of boxes → hundreds of per-figure CCIP
        calls + crops → a 30s+ job and an oversized submit (operator-flagged)."""
        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)[: self.cfg.max_figures]   # nms_merge is score-desc

    @staticmethod
    def _top(detector, image, cap: int):
        """Top-`cap` detections by score from an optional proposer (None → the
        proposer is off → []). Shared by the anatomy + panel proposers, which
        differ only in which detector and which cap."""
        if detector is None:
            return []
        return sorted(detector.detect(image), key=lambda b: b[1], reverse=True)[:cap]

    def components(self, image):
        return self._top(self._anatomy, image, self.cfg.max_components)

    def panels(self, image):
        return self._top(self._panel, image, self.cfg.max_panels)

    def unload(self) -> None:
        """Release every loaded proposer's YOLO (idle VRAM reclaim). The worker
        also drops its reference to this Proposers and rebuilds a fresh one via
        _proposers_for on the next job, so this is belt-and-braces."""
        for p in (self._person, self._anatomy, self._panel):
            if p is not None:
                p.unload()