miner.py

"""Scaffold miner for manak0/Detect-fire (specialized public package).

Required chute contract:
  - class named Miner
  - method predict_batch(batch_images, offset, n_keypoints) -> list[TVFrameResult]
  - this file lives at the root of the HF model repo

This scaffold is intentionally element-specialized (object labels,
element metadata). Weights are placeholder; distill/train fills real
ONNX/PT artifacts under the 30 MB hard cap.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from pydantic import BaseModel


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


ELEMENT_ID = 'manak0/Detect-fire'
SHORT_NAME = 'fire'
OBJECT_LABELS = (
    'fire',
    'smoke',
    'flame',
    'ember',
    'torch',
)


class Miner:
    """Specialist detector package for manak0/Detect-fire."""

    def __init__(self, path_hf_repo: Path) -> None:
        self.path_hf_repo = Path(path_hf_repo)
        self.element_id = ELEMENT_ID
        self.object_labels = list(OBJECT_LABELS)
        self._weights = self._discover_weights()

    def _discover_weights(self) -> Path | None:
        for name in ("model.onnx", "weights.onnx", "model.pt", "weights.pt"):
            cand = self.path_hf_repo / name
            if cand.is_file():
                return cand
        return None

    def __repr__(self) -> str:
        wname = self._weights.name if self._weights else None
        return (
            "Miner(element=%r, labels=%d, weights=%s)"
            % (self.element_id, len(self.object_labels), wname)
        )

    def predict_batch(
        self,
        batch_images: list[Any],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        """Return frame results. Scaffold emits empty boxes (schema-valid).

        Live distillation replaces this with a tiny specialist detector.
        """
        results: list[TVFrameResult] = []
        kps = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
        for i in range(len(batch_images)):
            results.append(
                TVFrameResult(
                    frame_id=offset + i,
                    boxes=[],
                    polygons=[],
                    keypoints=list(kps),
                )
            )
        return results