Upload folder using huggingface_hub

keypoint_helper.py

@@ -0,0 +1,116 @@
+
+import numpy as np
+from tqdm import tqdm
+from typing import List, Tuple, Sequence, Any
+
+FOOTBALL_KEYPOINTS: list[tuple[int, int]] = [
+    (0, 0),  # 1
+    (0, 0),  # 2
+    (0, 0),  # 3
+    (0, 0),  # 4
+    (0, 0),  # 5
+    (0, 0),  # 6
+    
+    (0, 0),  # 7
+    (0, 0),  # 8
+    (0, 0),  # 9
+
+    (0, 0),  # 10
+    (0, 0),  # 11
+    (0, 0),  # 12
+    (0, 0),  # 13
+
+    (0, 0),  # 14
+    (527, 283),  # 15
+    (527, 403),  # 16
+    (0, 0),  # 17
+
+    (0, 0),  # 18
+    (0, 0),  # 19
+    (0, 0),  # 20
+    (0, 0),  # 21
+
+    (0, 0),  # 22
+
+    (0, 0),  # 23
+    (0, 0),  # 24
+
+    (0, 0),  # 25
+    (0, 0),  # 26
+    (0, 0),  # 27
+    (0, 0),  # 28
+    (0, 0),  # 29
+    (0, 0),  # 30
+
+    (405, 340),  # 31
+    (645, 340),  # 32
+]
+
+def convert_keypoints_to_val_format(keypoints):
+    return [tuple(int(x) for x in pair) for pair in keypoints]
+
+def predict_failed_indices(results_frames: Sequence[Any]) -> list[int]:
+
+    max_frames = len(results_frames)
+    if max_frames == 0:
+        return []
+
+    failed_indices: list[int] = []
+    for frame_index, frame_result in enumerate(results_frames):
+        frame_keypoints = getattr(frame_result, "keypoints", []) or []
+        non_zero_count = sum(1 for (x, y) in frame_keypoints if int(x) != 0 and int(y) != 0)
+        if non_zero_count <= 4:
+            failed_indices.append(frame_index)
+    return failed_indices
+
+def _generate_sparse_template_keypoints(frame_width: int, frame_height: int) -> list[tuple[int, int]]:
+    template_max_x, template_max_y = (1045, 675)
+    sx = float(frame_width) / float(template_max_x if template_max_x != 0 else 1)
+    sy = float(frame_height) / float(template_max_y if template_max_y != 0 else 1)
+    scaled: list[tuple[int, int]] = []
+    for i in range(32):
+        tx, ty = FOOTBALL_KEYPOINTS[i]
+        x_scaled = int(round(tx * sx))
+        y_scaled = int(round(ty * sy))
+        scaled.append((x_scaled, y_scaled))
+    return scaled
+
+def fix_keypoints(
+    results_frames: Sequence[Any],
+    failed_indices: Sequence[int],
+    frame_width: int,
+    frame_height: int,
+) -> list[Any]:
+    max_frames = len(results_frames)
+    if max_frames == 0:
+        return list(results_frames)
+
+    failed_set = set(int(i) for i in failed_indices)
+    all_indices = list(range(max_frames))
+    successful_indices = [i for i in all_indices if i not in failed_set]
+
+    if len(successful_indices) == 0:
+        sparse_template = _generate_sparse_template_keypoints(frame_width, frame_height)
+        for frame_result in results_frames:
+            setattr(frame_result, "keypoints", list(convert_keypoints_to_val_format(sparse_template)))
+        return list(results_frames)
+
+    seed_index = successful_indices[0]
+    seed_kps_raw = getattr(results_frames[seed_index], "keypoints", []) or []
+    last_success_kps = convert_keypoints_to_val_format(seed_kps_raw)
+
+    for frame_index in range(max_frames):
+        frame_result = results_frames[frame_index]
+        if frame_index in failed_set:
+            setattr(frame_result, "keypoints", list(last_success_kps))
+        else:
+            current_kps_raw = getattr(frame_result, "keypoints", []) or []
+            current_kps = convert_keypoints_to_val_format(current_kps_raw)
+            setattr(frame_result, "keypoints", list(current_kps))
+            last_success_kps = current_kps
+
+    return list(results_frames)
+
+def run_keypoints_post_processing(results_frames: Sequence[Any], frame_width: int, frame_height: int) -> list[Any]:
+    failed_indices = predict_failed_indices(results_frames)
+    return fix_keypoints(results_frames, failed_indices, frame_width, frame_height)

miner.py

@@ -6,7 +6,7 @@ import os
 from numpy import ndarray
 from pydantic import BaseModel
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
+from keypoint_helper import run_keypoints_post_processing
 
 from ultralytics import YOLO
 from team_cluster import TeamClassifier
@@ -46,7 +46,7 @@ class Miner:
     CORNER_CONFIDENCE = Constants.CORNER_CONFIDENCE
     GOALKEEPER_POSITION_MARGIN = Constants.GOALKEEPER_POSITION_MARGIN
     MIN_SAMPLES_FOR_FIT = 16  # Minimum player crops needed before fitting TeamClassifier
-    MAX_SAMPLES_FOR_FIT = 1000  # Maximum samples to avoid overfitting
+    MAX_SAMPLES_FOR_FIT = 600  # Maximum samples to avoid overfitting
 
     def __init__(self, path_hf_repo: Path) -> None:
         try:
@@ -54,7 +54,7 @@ class Miner:
             model_path = path_hf_repo / "football_object_detection.onnx"
             self.bbox_model = YOLO(model_path)
             
-            print("✅ BBox Model Loaded")
+            print("BBox Model Loaded")
 
             team_model_path = path_hf_repo / "osnet_model.pth.tar-100"
             self.team_classifier = TeamClassifier(
@@ -62,7 +62,7 @@ class Miner:
                 batch_size=32,
                 model_name=str(team_model_path)
             )
-            print("✅ Team Classifier Loaded")
+            print("Team Classifier Loaded")
             
             # Team classification state
             self.team_classifier_fitted = False
@@ -146,7 +146,7 @@ class Miner:
         self.team_classifier_fitted = False
         start = time.time()
         # Collect player crops from first batch for fitting
-        fit_sample_size = 1000
+        fit_sample_size = 600
         player_crops_for_fit = []
 
         for frame_id in range(len(detection_results)):
@@ -422,6 +422,10 @@ class Miner:
             )
             results.append(result)
 
+        if len(batch_images) > 0:
+            h, w = batch_images[0].shape[:2]
+            results = run_keypoints_post_processing(results, w, h)
+
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()