Update miner.py

miner.py

@@ -1,53 +1,34 @@
 from pathlib import Path
 import math
+
 import cv2
 import numpy as np
 import onnxruntime as ort
 from numpy import ndarray
-import struct
-import shutil
-import importlib.util as imutil
-import sys
-
-class BoundingBox:
-    def __init__(self, x1: int, y1: int, x2: int, y2: int,
-                 cls_id: int, conf: float) -> None:
-        self.x1 = x1
-        self.y1 = y1
-        self.x2 = x2
-        self.y2 = y2
-        self.cls_id = cls_id
-        self.conf = conf
-
-    def model_dump(self, *args, **kwargs) -> dict:
-        return {
-            "x1": self.x1,
-            "y1": self.y1,
-            "x2": self.x2,
-            "y2": self.y2,
-            "cls_id": self.cls_id,
-            "conf": self.conf,
-        }
-
-
-class TVFrameResult:
-    def __init__(self, frame_id: int, boxes: list[BoundingBox],
-                 keypoints: list[tuple[int, int]]) -> None:
-        self.frame_id = frame_id
-        self.boxes = boxes
-        self.keypoints = keypoints
-
-    def model_dump(self, *args, **kwargs) -> dict:
-        return {
-            "frame_id": self.frame_id,
-            "boxes": [box.model_dump() for box in self.boxes],
-            "keypoints": self.keypoints,
-        }
+from pydantic import BaseModel
+
+
+class BoundingBox(BaseModel):
+    x1: int
+    y1: int
+    x2: int
+    y2: int
+    cls_id: int
+    conf: float
+
+
+class TVFrameResult(BaseModel):
+    frame_id: int
+    boxes: list[BoundingBox]
+    keypoints: list[tuple[int, int]]
 
 
 class Miner:
     """ONNX Runtime miner. Hard global NMS + sanity filter + dedup + flip TTA, with per-class rescue bonus."""
+
     class_names = ["cup", "bottle", "can"]
+    model_class_names = ["bottle", "can", "cup"]
+    _model_to_competition_cls = np.array([1, 2, 0], dtype=np.int32)
     input_size = 1280
     iou_thres = 0.4
     cross_iou_thresh = 0.7
@@ -57,44 +38,9 @@ class Miner:
     max_det = 300
     _conf_thres_array = np.array([0.6, 0.45, 0.5], dtype=np.float32)
     _bonus_array = np.array([0.0, 0.0, 0.2], dtype=np.float32)
-    
-    def load_model_from_onnx(self, model_path, model_name) -> None:
-        so_path = f"{model_name}.so"
-        shutil.copy2(model_path, so_path)
-        spec = imutil.spec_from_file_location(
-            model_name,
-            so_path
-        )
-        model = imutil.module_from_spec(spec)
-        spec.loader.exec_module(model)
-        sys.modules[model_name] = model
-        self.model = model
 
     def __init__(self, path_hf_repo: Path) -> None:
-        model_path = "./weight.onnx"
-        self.class_names = ["cup", "bottle", "can"]
-        with open(path_hf_repo / "weights.onnx", "rb") as f:
-            f.seek(-16, 2)
-            so_size = struct.unpack("<Q", f.read(8))[0]
-            model_size = struct.unpack("<Q", f.read(8))[0]
-            f.seek(0)
-            model = f.read(so_size)
-            weight = f.read(model_size)
-        with open("model.onnx", "wb") as f:
-            f.write(model)
-        with open("weight.onnx", "wb") as f:
-            f.write(weight)
-
-        self.model = None
-        try:
-            self.load_model_from_onnx("model.onnx", "model")
-        except Exception as e:
-            print(
-                "Embedded model module failed to load; "
-                f"falling back to ONNX Runtime inference: {type(e).__name__}: {e}"
-            )
-        
-        self.class_names = ["cup", "bottle", "can"]
+        model_path = path_hf_repo / "weights.onnx"
         print("ORT version:", ort.__version__)
 
         try:
@@ -133,6 +79,7 @@ class Miner:
         self.input_name = self.session.get_inputs()[0].name
         self.output_names = [output.name for output in self.session.get_outputs()]
         self.input_shape = self.session.get_inputs()[0].shape
+        self.input_dtype = self._input_dtype(self.session.get_inputs()[0].type)
 
         self.input_height = self._safe_dim(self.input_shape[2], default=self.input_size)
         self.input_width = self._safe_dim(self.input_shape[3], default=self.input_size)
@@ -140,6 +87,7 @@ class Miner:
         print(f"ONNX model loaded from: {model_path}")
         print(f"ONNX providers: {self.session.get_providers()}")
         print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
+        print(f"ONNX input dtype: {self.input_dtype}")
 
     def __repr__(self) -> str:
         return (
@@ -151,6 +99,21 @@ class Miner:
     def _safe_dim(value, default: int) -> int:
         return value if isinstance(value, int) and value > 0 else default
 
+    @staticmethod
+    def _input_dtype(input_type: str) -> np.dtype:
+        if input_type == "tensor(float16)":
+            return np.dtype(np.float16)
+        return np.dtype(np.float32)
+
+    @classmethod
+    def _to_competition_cls(cls, cls_ids: np.ndarray) -> np.ndarray:
+        if len(cls_ids) == 0:
+            return cls_ids.astype(np.int32)
+        valid = (cls_ids >= 0) & (cls_ids < len(cls._model_to_competition_cls))
+        remapped = np.full_like(cls_ids, fill_value=-1, dtype=np.int32)
+        remapped[valid] = cls._model_to_competition_cls[cls_ids[valid]]
+        return remapped
+
     def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
                    color=(114, 114, 114)
                    ) -> tuple[ndarray, float, tuple[float, float]]:
@@ -178,9 +141,9 @@ class Miner:
         orig_h, orig_w = image.shape[:2]
         img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        img = img.astype(np.float32) / 255.0
+        img = img.astype(self.input_dtype) / self.input_dtype.type(255.0)
         img = np.transpose(img, (2, 0, 1))[None, ...]
-        img = np.ascontiguousarray(img, dtype=np.float32)
+        img = np.ascontiguousarray(img, dtype=self.input_dtype)
         return img, ratio, pad, (orig_w, orig_h)
 
     @staticmethod
@@ -220,9 +183,9 @@ class Miner:
             xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
             yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
             inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
-            a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) * 
+            a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
                    max(0.0, boxes[i, 3] - boxes[i, 1]))
-            a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) * 
+            a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
                    np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
             iou = inter / (a_i + a_r - inter + 1e-7)
             order = rest[iou <= iou_thresh]
@@ -251,7 +214,7 @@ class Miner:
         boxes = np.asarray(boxes, dtype=np.float32)
         scores = np.asarray(scores, dtype=np.float32)
         cls_ids = np.asarray(cls_ids, dtype=np.int32)
-        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * 
+        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
                  np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
         margins = scores - self._conf_thres_array[cls_ids]
         order = np.lexsort((-areas, -margins))
@@ -291,9 +254,9 @@ class Miner:
             np.inf,
         )
         keep = (
-            (bw >= self.min_side) & (bh >= self.min_side) & 
-            (area >= self.min_box_area) & 
-            (area <= 0.95 * image_area) & 
+            (bw >= self.min_side) & (bh >= self.min_side) &
+            (area >= self.min_box_area) &
+            (area <= 0.95 * image_area) &
             (ar <= self.max_aspect_ratio)
         )
         return boxes[keep], scores[keep], cls_ids[keep]
@@ -307,7 +270,7 @@ class Miner:
         n = len(post_boxes)
         if n == 0:
             return np.empty(0, dtype=np.float32)
-        full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) * 
+        full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
                       np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
         out = np.empty(n, dtype=np.float32)
         for i in range(n):
@@ -373,9 +336,13 @@ class Miner:
         if preds.ndim != 2 or preds.shape[1] < 6:
             raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
 
-        boxes = preds[:,:4].astype(np.float32)
+        boxes = preds[:, :4].astype(np.float32)
         scores = preds[:, 4].astype(np.float32)
-        cls_ids = preds[:, 5].astype(np.int32)
+        cls_ids = self._to_competition_cls(preds[:, 5].astype(np.int32))
+        valid_cls = cls_ids >= 0
+        boxes = boxes[valid_cls]
+        scores = scores[valid_cls]
+        cls_ids = cls_ids[valid_cls]
 
         keep = self._conf_filter_mask(scores, cls_ids)
         boxes = boxes[keep]
@@ -406,7 +373,7 @@ class Miner:
         if preds.ndim != 2 or preds.shape[1] < 5:
             raise ValueError(f"Unexpected raw output shape: {preds.shape}")
 
-        boxes_xywh = preds[:,:4].astype(np.float32)
+        boxes_xywh = preds[:, :4].astype(np.float32)
         cls_part = preds[:, 4:].astype(np.float32)
         if cls_part.shape[1] == 1:
             scores = cls_part[:, 0]
@@ -414,6 +381,11 @@ class Miner:
         else:
             cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
             scores = cls_part[np.arange(len(cls_part)), cls_ids]
+        cls_ids = self._to_competition_cls(cls_ids)
+        valid_cls = cls_ids >= 0
+        boxes_xywh = boxes_xywh[valid_cls]
+        scores = scores[valid_cls]
+        cls_ids = cls_ids[valid_cls]
 
         keep = self._conf_filter_mask(scores, cls_ids)
         boxes_xywh = boxes_xywh[keep]
@@ -486,81 +458,61 @@ class Miner:
         return self._postprocess(outputs[0], ratio, pad, orig_size)
 
     def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
-        if self.model is None:
-            boxes_orig = self._predict_single(image)
-            flipped = cv2.flip(image, 1)
-            boxes_flip = self._predict_single(flipped)
-            w = image.shape[1]
-            boxes_flip = [
-                BoundingBox(
-                    x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
-                    cls_id=b.cls_id, conf=b.conf,
-                )
-                for b in boxes_flip
-            ]
-            all_boxes = boxes_orig + boxes_flip
-            if not all_boxes:
-                return []
-            coords = np.array(
-                [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
+        boxes_orig = self._predict_single(image)
+        flipped = cv2.flip(image, 1)
+        boxes_flip = self._predict_single(flipped)
+        w = image.shape[1]
+        boxes_flip = [
+            BoundingBox(
+                x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
+                cls_id=b.cls_id, conf=b.conf,
             )
-            scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
-            cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
-
-            hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
-            if len(hard_keep) == 0:
-                return []
-            if len(hard_keep) > self.max_det:
-                top = np.argsort(-scores[hard_keep])[: self.max_det]
-                hard_keep = hard_keep[top]
-            boosted = self._max_score_per_cluster(
-                coords[hard_keep], cls_ids[hard_keep],
-                coords, scores, cls_ids, self.iou_thres,
+            for b in boxes_flip
+        ]
+        all_boxes = boxes_orig + boxes_flip
+        if not all_boxes:
+            return []
+
+        coords = np.array(
+            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
+        )
+        scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
+        cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
+
+        hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
+        if len(hard_keep) == 0:
+            return []
+        if len(hard_keep) > self.max_det:
+            top = np.argsort(-scores[hard_keep])[: self.max_det]
+            hard_keep = hard_keep[top]
+        boosted = self._max_score_per_cluster(
+            coords[hard_keep], cls_ids[hard_keep],
+            coords, scores, cls_ids, self.iou_thres,
+        )
+
+        kept_coords = coords[hard_keep]
+        kept_cls = cls_ids[hard_keep]
+        if len(kept_coords) > 1:
+            kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
+                kept_coords, boosted, kept_cls, self.cross_iou_thresh
             )
-            kept_coords = coords[hard_keep]
-            kept_cls = cls_ids[hard_keep]
-            if len(kept_coords) > 1:
-                kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
-                    kept_coords, boosted, kept_cls, self.cross_iou_thresh
-                )
-            return [
-                BoundingBox(
-                    x1=int(math.floor(kept_coords[j, 0])),
-                    y1=int(math.floor(kept_coords[j, 1])),
-                    x2=int(math.ceil(kept_coords[j, 2])),
-                    y2=int(math.ceil(kept_coords[j, 3])),
-                    cls_id=int(kept_cls[j]),
-                    conf=float(boosted[j]),
-                )
-                for j in range(len(kept_coords))
-            ]
-        else:
-            output = self.model.predict(image)
-            print(output)
-            res: list[BoundingBox] = []
-            for each_output in output["results"]:
-                cls_id = self.class_names.index(each_output["echo"]["text"])
-                for prediction in each_output["predictions"]:
-                    for polygon in prediction["masks"]:
-                        pts = np.array(polygon, dtype=np.int32)
-                        x, y, w, h = cv2.boundingRect(pts)
-                        res.append(
-                            BoundingBox(
-                                x1=int(x),
-                                y1=int(y),
-                                x2=int(x + w),
-                                y2=int(y + h),
-                                cls_id=cls_id,
-                                conf=float(prediction["confidence"]),
-                            )
-                        )
-            return res
+
+        return [
+            BoundingBox(
+                x1=int(math.floor(kept_coords[j, 0])),
+                y1=int(math.floor(kept_coords[j, 1])),
+                x2=int(math.ceil(kept_coords[j, 2])),
+                y2=int(math.ceil(kept_coords[j, 3])),
+                cls_id=int(kept_cls[j]),
+                conf=float(boosted[j]),
+            )
+            for j in range(len(kept_coords))
+        ]
 
     def predict_batch(self, batch_images: list[ndarray], offset: int,
                       n_keypoints: int) -> list[TVFrameResult]:
         results: list[TVFrameResult] = []
         for frame_number_in_batch, image in enumerate(batch_images):
-            print(image.shape)
             try:
                 boxes = self._predict_tta(image)
             except Exception as e:
@@ -573,4 +525,4 @@ class Miner:
                     keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
                 )
             )
-        return results
+        return results