scorevision: push artifact

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)