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
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]