change object detection model

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae707480215c937e905a7236c1f72867231e8e475eab4e91bebd252aa64236d8
+size 5455642

miner.py

@@ -61,10 +61,18 @@ class Miner:
     SMALL_RATIO_MAX: float = 0.50
     SINGLE_PLAYER_HUE_PIVOT: float = 90.0
 
+    CLS_MAP = {
+        0: 0,
+        1: 1,
+        2: 6,
+        3: 7,
+        4: 3,
+    }
+
     def __init__(self, path_hf_repo: Path) -> None:
-        print(path_hf_repo / "objdetect.pt")
-        self.bbox_model = YOLO(path_hf_repo / "objdetect.pt")
-        print(" BBox Model (objdetect.pt) Loaded")
+        print(path_hf_repo / "best.pt")
+        self.bbox_model = YOLO(path_hf_repo / "best.pt")
+        print(" BBox Model (best.pt) Loaded")
         device = "cuda" if torch.cuda.is_available() else "cpu"
         # model_kp_path = path_hf_repo / "SV_kp.engine"
         # model_kp = torch_tensorrt.load(model_kp_path)
@@ -255,7 +263,7 @@ class Miner:
                             y1=int(y1),
                             x2=int(x2),
                             y2=int(y2),
-                            cls_id=int(cls_id),
+                            cls_id=int(self.CLS_MAP[int(cls_id)]),
                             conf=float(conf),
                         )
                     )
@@ -264,32 +272,32 @@ class Miner:
                     best_ball = max(footballs, key=lambda b: b.conf)
                     boxes = [bb for bb in boxes if int(bb.cls_id) != 0]
                     boxes.append(best_ball)
-                boxes = self.suppress_quasi_total_containment(boxes)
-                boxes = self.suppress_small_contained(boxes)
-                img_bgr = batch_images[frame_idx_in_batch]
-                player_indices: List[int] = []
-                player_feats: List[np.ndarray] = []
-                for i, bb in enumerate(boxes):
-                    if int(bb.cls_id) == 2:
-                        hs = self._hs_feature_from_roi(img_bgr, bb)
-                        player_indices.append(i)
-                        player_feats.append(hs)
-                cluster_centers = None
-                n_players = len(player_feats)
-                if n_players >= 2:
-                    feats = np.vstack(player_feats)
-                    labels, centers = self._assign_players_two_clusters(feats)
-                    order = np.argsort(centers[:, 0])
-                    centers = centers[order]
-                    remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)}
-                    labels = np.vectorize(remap.get)(labels)
-                    cluster_centers = centers
-                    for idx_in_list, lbl in zip(player_indices, labels):
-                        boxes[idx_in_list].cls_id = 6 if int(lbl) == 0 else 7
-                elif n_players == 1:
-                    hue, _ = player_feats[0]
-                    boxes[player_indices[0]].cls_id = 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7
-                self._reclass_extra_goalkeepers(img_bgr, boxes, cluster_centers)
+                # boxes = self.suppress_quasi_total_containment(boxes)
+                # boxes = self.suppress_small_contained(boxes)
+                # img_bgr = batch_images[frame_idx_in_batch]
+                # player_indices: List[int] = []
+                # player_feats: List[np.ndarray] = []
+                # for i, bb in enumerate(boxes):
+                #     if int(bb.cls_id) == 2:
+                #         hs = self._hs_feature_from_roi(img_bgr, bb)
+                #         player_indices.append(i)
+                #         player_feats.append(hs)
+                # cluster_centers = None
+                # n_players = len(player_feats)
+                # if n_players >= 2:
+                #     feats = np.vstack(player_feats)
+                #     labels, centers = self._assign_players_two_clusters(feats)
+                #     order = np.argsort(centers[:, 0])
+                #     centers = centers[order]
+                #     remap = {old_idx: new_idx for new_idx, old_idx in enumerate(order)}
+                #     labels = np.vectorize(remap.get)(labels)
+                #     cluster_centers = centers
+                #     for idx_in_list, lbl in zip(player_indices, labels):
+                #         boxes[idx_in_list].cls_id = 6 if int(lbl) == 0 else 7
+                # elif n_players == 1:
+                #     hue, _ = player_feats[0]
+                #     boxes[player_indices[0]].cls_id = 6 if float(hue) < self.SINGLE_PLAYER_HUE_PIVOT else 7
+                # self._reclass_extra_goalkeepers(img_bgr, boxes, cluster_centers)
                 bboxes[offset + frame_idx_in_batch] = boxes
 
         pitch_batch_size = min(self.pitch_batch_size, len(batch_images))