|
|
| import numpy as np
|
| from tqdm import tqdm
|
| from typing import List, Tuple, Sequence, Any
|
|
|
| FOOTBALL_KEYPOINTS: list[tuple[int, int]] = [
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
|
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
|
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
|
|
| (0, 0),
|
| (527, 283),
|
| (527, 403),
|
| (0, 0),
|
|
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
|
|
| (0, 0),
|
|
|
| (0, 0),
|
| (0, 0),
|
|
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
| (0, 0),
|
|
|
| (405, 340),
|
| (645, 340),
|
| ]
|
|
|
| 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) |
