miner.py

"""SN44 public-track miner entrypoint. Goes at the ROOT of your HF repo.

Sandbox constraints (verified against
scorevision/validator/audit/open_source/security.py):
  * ALL logic must live in THIS file - the chute installs an import blocker
    that rejects modules loaded outside stdlib/site-packages.
  * Class `Miner` with `predict_batch(batch_images, offset, n_keypoints)`;
    parameter NAMES are checked by signature inspection.
  * Banned imports: socket, subprocess, ctypes, multiprocessing, requests,
    urllib, http, ftplib, telnetlib, paramiko.
    Banned calls: eval, exec, __import__, open, os.system/popen/remove/...
  * Model artifacts: .onnx ONLY. Repo <= 30 MB.
  * Single-frame p95 <= 110 ms on ~4 CPU threads.

CLASS ORDER IS THE #1 SILENT KILLER. The validator maps a prediction's
cls_id through the manifest `objects` list:
    manak0/Detect-fire -> ["fire", "smoke", "fire extinguisher"]
Ultralytics models are commonly exported with a DIFFERENT internal order
(Score's own reference uses [fire, fire extinguisher, smoke]). We read the
`names` metadata Ultralytics embeds in the ONNX and remap onto manifest
order at runtime, so a retrained model with a different order still works.
An out-of-range or mis-mapped cls_id is dropped silently by the validator -
indistinguishable from a broken model.

SCORING (measured on 157 real challenges of the incumbent):
    raw = 0.6*map50 + 0.4*false_positive
    false_positive = max(0, 1 - (total_FP / n_images)/10)
Predicting nothing already scores raw 0.40, so a loose threshold is
expensive. Tune with miner_dev.sweep against the real metric, not mAP.
"""

import ast
import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
from pydantic import BaseModel

MANIFEST_OBJECTS = ["fire", "smoke", "fire extinguisher"]


# predict_batch MUST return objects exposing .model_dump(), not plain dicts.
# The live chute does `fr = frame_result.model_dump()` unconditionally
# (chute_template/turbovision_chute.py.j2), while the compliance runner does
# `model_dump() if hasattr(...) else dict(frame_result)`. Returning dicts
# therefore PASSES compliance and fails every real challenge with
# "'dict' object has no attribute 'model_dump'" - a successful HTTP 200 whose
# body is {"success": false}, scored as zero. Mirror the reference contract in
# scorevision/miner/open_source/example_miner/miner.py exactly.
class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class Polygon(BaseModel):
    cls_id: int
    conf: float
    points: list[tuple[int, int]]


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list[Polygon] | None = None
    keypoints: list[tuple[int, int]] | None = None

MODEL_FILE = "model.onnx"
# 640/672/768 all fit the CPU budget; the incumbent runs 672 at 43.7 ms p95
# on a 4-thread box against a 110 ms ceiling, so 768 is affordable and buys
# map50 on small/distant objects - which is where the headroom is.
INPUT_SIZE = 704
NUM_THREADS = 4

# Per-class confidence thresholds, indexed by MANIFEST order.
CONF_THRES = np.array([0.2, 0.2, 0.15], dtype=np.float32)
# If a class has ZERO boxes over threshold, admit its top-1 candidate when it
# scores at least (threshold - bonus). Recovers recall on borderline frames
# without paying the false-positive cost on frames that already have boxes.
RESCUE_BONUS = np.array([0.03, 0.10, 0.05], dtype=np.float32)

IOU_THRES = 0.55          # per-class NMS (only used for non-end2end heads)
SAME_IOU_THRES = 0.70     # same-class dedup; end2end o2o heads still emit near-duplicates
CROSS_IOU_THRES = 0.90    # cross-class duplicate suppression, by IoU (see _postprocess)
MAX_DET = 30

# Box sanity filter: drop degenerate / tiny / image-spanning detections.
MIN_BOX_AREA = 14 * 14
MIN_SIDE = 8
MAX_ASPECT = 8.0
MAX_AREA_FRAC = 0.92

# Same-class union-merge when intersection covers this fraction of the
# SMALLER box. Smoke plumes fragment, so merging helps; separate flames must
# stay separate, so fire is disabled (>1.0).
MERGE_OVERLAP = np.array([1.01, 0.80, 1.01], dtype=np.float32)


def _letterbox(img, size):
    import cv2
    h, w = img.shape[:2]
    s = min(size / max(h, 1), size / max(w, 1))
    nh, nw = max(1, int(round(h * s))), max(1, int(round(w * s)))
    canvas = np.full((size, size, 3), 114, dtype=np.uint8)
    dy, dx = (size - nh) // 2, (size - nw) // 2
    canvas[dy:dy + nh, dx:dx + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
    return canvas, s, dx, dy


def _nms(boxes, scores, thr):
    if boxes.size == 0:
        return []
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size:
        i = order[0]
        keep.append(int(i))
        if order.size == 1:
            break
        xx1 = np.maximum(x1[i], x1[order[1:]]); yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]]); yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        union = areas[i] + areas[order[1:]] - inter
        # np.where would evaluate eagerly and emit nan on union==0; nan <= thr
        # is False, which would silently DROP a valid box.
        iou = np.divide(inter, union, out=np.zeros_like(inter, dtype=np.float64), where=union > 0)
        order = order[1:][iou <= thr]
    return keep


def _inter_over_smaller(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    sa = max(0.0, a[2] - a[0]) * max(0.0, a[3] - a[1])
    sb = max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1])
    m = min(sa, sb)
    return inter / m if m > 0 else 0.0


def _iou_pair(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    ua = (max(0.0, a[2] - a[0]) * max(0.0, a[3] - a[1])
          + max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1]) - inter)
    return inter / ua if ua > 0 else 0.0


class Miner:
    def __init__(self, path_hf_repo) -> None:
        repo = Path(path_hf_repo)
        model_path = repo / MODEL_FILE
        if not model_path.is_file():
            raise FileNotFoundError(f"missing {MODEL_FILE} in {repo}")

        opts = ort.SessionOptions()
        opts.intra_op_num_threads = NUM_THREADS
        opts.inter_op_num_threads = 1
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(str(model_path), opts,
                                            providers=["CPUExecutionProvider"])
        inp = self.session.get_inputs()[0]
        self.input_name = inp.name
        # The exported model's own spatial size is authoritative. Forcing a
        # different INPUT_SIZE against a static-shape export raises, and the
        # never-raise handler in predict_batch would turn that into a silent
        # zero score. Trust the graph; fall back to INPUT_SIZE only if dynamic.
        static = [d for d in inp.shape[2:] if isinstance(d, int) and d > 0]
        self.size = int(static[0]) if len(static) == 2 else INPUT_SIZE
        self.remap = self._build_remap()
        self.end2end = None  # resolved on first inference from output shape
        self.last_error = None  # surfaced for debugging; never raised

    def _build_remap(self):
        """Model class index -> manifest index, by NAME, from ONNX metadata."""
        try:
            meta = self.session.get_modelmeta().custom_metadata_map or {}
            raw = meta.get("names")
            names = ast.literal_eval(raw) if raw else None
            if isinstance(names, dict):
                out = {}
                for k, v in names.items():
                    n = str(v).strip().lower()
                    if n in MANIFEST_OBJECTS:
                        out[int(k)] = MANIFEST_OBJECTS.index(n)
                if out:
                    return out
        except Exception:
            pass
        return {i: i for i in range(len(MANIFEST_OBJECTS))}

    def __repr__(self):
        return (f"ONNX detector size={self.size} threads={NUM_THREADS} "
                f"remap={self.remap} conf={CONF_THRES.tolist()}")

    def _decode(self, raw, s, dx, dy, h, w):
        """Return (xyxy Nx4, cls N, conf N) in ORIGINAL image coords."""
        arr = raw[0] if raw.ndim == 3 else raw
        if arr.ndim == 2 and arr.shape[-1] == 6:      # end2end: already NMS'd
            self.end2end = True
            boxes = arr[:, :4].astype(np.float32)
            conf = arr[:, 4].astype(np.float32)
            cls = arr[:, 5].astype(np.int32)
            keep = conf > 0
            boxes, conf, cls = boxes[keep], conf[keep], cls[keep]
        else:                                          # raw head -> needs NMS
            self.end2end = False
            pred = arr.T if arr.shape[0] < arr.shape[1] else arr
            if pred.shape[1] < 5:
                return np.zeros((0, 4)), np.zeros(0, int), np.zeros(0)
            xywh, sc = pred[:, :4], pred[:, 4:]
            cls = sc.argmax(1).astype(np.int32)
            conf = sc.max(1).astype(np.float32)
            keep = conf >= float(CONF_THRES.min() - RESCUE_BONUS.max())
            xywh, cls, conf = xywh[keep], cls[keep], conf[keep]
            cx, cy, bw, bh = xywh[:, 0], xywh[:, 1], xywh[:, 2], xywh[:, 3]
            boxes = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], 1)
            sel = []
            for c in np.unique(cls):
                m = np.where(cls == c)[0]
                sel.extend(m[_nms(boxes[m], conf[m], IOU_THRES)])
            sel = np.array(sorted(sel), dtype=int) if sel else np.zeros(0, int)
            boxes, cls, conf = boxes[sel], cls[sel], conf[sel]

        if boxes.shape[0]:
            boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dx) / max(s, 1e-9)
            boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dy) / max(s, 1e-9)
            boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
            boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
        return boxes, cls, conf

    def _postprocess(self, boxes, cls, conf, h, w):
        # remap model classes onto manifest order, drop unknown classes
        mapped = np.array([self.remap.get(int(c), -1) for c in cls], dtype=np.int32)
        ok = mapped >= 0
        boxes, conf, mapped = boxes[ok], conf[ok], mapped[ok]
        if not boxes.shape[0]:
            return []

        # box sanity filter
        bw = boxes[:, 2] - boxes[:, 0]
        bh = boxes[:, 3] - boxes[:, 1]
        area = bw * bh
        with np.errstate(divide="ignore", invalid="ignore"):
            ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
        sane = ((bw >= MIN_SIDE) & (bh >= MIN_SIDE) & (area >= MIN_BOX_AREA)
                & (ar <= MAX_ASPECT) & (area <= MAX_AREA_FRAC * h * w))
        boxes, conf, mapped = boxes[sane], conf[sane], mapped[sane]
        if not boxes.shape[0]:
            return []

        # per-class threshold + rescue bonus
        keep_idx = []
        for c in range(len(MANIFEST_OBJECTS)):
            m = np.where(mapped == c)[0]
            if not m.size:
                continue
            passing = m[conf[m] >= CONF_THRES[c]]
            if passing.size:
                keep_idx.extend(passing.tolist())
            else:
                top = m[int(np.argmax(conf[m]))]
                if conf[top] >= CONF_THRES[c] - RESCUE_BONUS[c]:
                    keep_idx.append(int(top))
        if not keep_idx:
            return []
        keep_idx = np.array(sorted(set(keep_idx)), dtype=int)
        boxes, conf, mapped = boxes[keep_idx], conf[keep_idx], mapped[keep_idx]

        # same-class dedup. The end2end branch skips NMS entirely, but the o2o head
        # still emits near-duplicates; each one is scored as a false positive AND
        # steals no match, so it is pure loss under the adaptive-IoU rule.
        sel = []
        for c in range(len(MANIFEST_OBJECTS)):
            m = np.where(mapped == c)[0]
            if m.size:
                sel.extend(m[_nms(boxes[m], conf[m], SAME_IOU_THRES)])
        if not sel:
            return []
        sel = np.array(sorted(sel), dtype=int)
        boxes, conf, mapped = boxes[sel], conf[sel], mapped[sel]

        # same-class union merge (smoke fragments; fire disabled)
        for c in range(len(MANIFEST_OBJECTS)):
            if MERGE_OVERLAP[c] > 1.0:
                continue
            changed = True
            while changed:
                changed = False
                idx = np.where(mapped == c)[0]
                for a in range(len(idx)):
                    for b in range(a + 1, len(idx)):
                        i, j = idx[a], idx[b]
                        if _inter_over_smaller(boxes[i], boxes[j]) >= MERGE_OVERLAP[c]:
                            boxes[i] = [min(boxes[i][0], boxes[j][0]), min(boxes[i][1], boxes[j][1]),
                                        max(boxes[i][2], boxes[j][2]), max(boxes[i][3], boxes[j][3])]
                            conf[i] = max(conf[i], conf[j])
                            mapped[j] = -1
                            changed = True
                            break
                    if changed:
                        break
                sel = mapped >= 0
                boxes, conf, mapped = boxes[sel], conf[sel], mapped[sel]

        # Cross-class duplicate suppression: same object carrying two labels.
        # Must be IoU, not intersection-over-smaller: fire sits *inside* smoke in
        # most real frames, which drives IoS to ~1.0 and deleted the true fire box.
        order = conf.argsort()[::-1]
        dead = set()
        for a in range(len(order)):
            i = order[a]
            if i in dead:
                continue
            for b in range(a + 1, len(order)):
                j = order[b]
                if j in dead or mapped[i] == mapped[j]:
                    continue
                if _iou_pair(boxes[i], boxes[j]) >= CROSS_IOU_THRES:
                    dead.add(j)

        out = []
        for i in order:
            if i in dead:
                continue
            x1, y1, x2, y2 = boxes[i]
            if x2 <= x1 or y2 <= y1:
                continue
            out.append(BoundingBox(x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
                                   cls_id=int(mapped[i]), conf=float(conf[i])))
            if len(out) >= MAX_DET:
                break
        return out

    def predict_batch(self, batch_images, offset: int, n_keypoints: int) -> list:
        """Signature is contract-checked: do not rename these parameters."""
        results = []
        for i, img in enumerate(batch_images):
            frame_id = offset + i
            try:
                arr = np.asarray(img)
                if arr.ndim == 2:
                    arr = np.stack([arr] * 3, axis=-1)
                h, w = arr.shape[:2]
                canvas, s, dx, dy = _letterbox(arr, self.size)
                blob = canvas[:, :, ::-1].transpose(2, 0, 1)[None].astype(np.float32) / 255.0
                raw = self.session.run(None, {self.input_name: blob})[0]
                boxes, cls, conf = self._decode(raw, s, dx, dy, h, w)
                dets = self._postprocess(boxes, cls, conf, h, w)
            except Exception as e:
                # Never raise: an exception zeroes the whole challenge. Record
                # it so offline harnesses can tell "no detections" apart from
                # "crashed" - silent except made those indistinguishable.
                self.last_error = f"{type(e).__name__}: {e}"
                dets = []
            results.append(TVFrameResult(frame_id=frame_id, boxes=dets,
                                         polygons=[], keypoints=[]))
        return results