| 73 | |
| 74 | |
| 75 | class PlayerTracker: |
| 76 | def __init__(self, model_path='yolov8m-pose.pt', device='cuda', imgsz: int = 1280, conf: float = 0.05): |
| 77 | """ |
| 78 | Initialize player tracker with a YOLOv8-Pose model. |
| 79 | |
| 80 | Using a pose model allows extracting 17 body keypoints per player |
| 81 | per frame in the same inference pass — no extra compute cost. |
| 82 | |
| 83 | Args: |
| 84 | model_path: Path to the YOLOv8-Pose model (e.g. 'yolov8m-pose.pt'). |
| 85 | device: 'cuda' or 'cpu' |
| 86 | imgsz: YOLO inference resolution. Should match input video resolution |
| 87 | to avoid downscaling small far-player detections below threshold. |
| 88 | conf: Detection confidence threshold. Lower values surface more |
| 89 | candidates (needed for small far-court players). |
| 90 | """ |
| 91 | self.model = YOLO(model_path) |
| 92 | self.imgsz = imgsz |
| 93 | self.conf = conf |
| 94 | if device == 'cuda': |
| 95 | self.model.to(device) |
| 96 | |
| 97 | def detect_frame(self, frame) -> tuple[dict, dict]: |
| 98 | """ |
| 99 | Detect players in a single frame and extract pose keypoints. |
| 100 | |
| 101 | Args: |
| 102 | frame: Single video frame (numpy array, BGR). |
| 103 | |
| 104 | Returns: |
| 105 | player_dict: {track_id: [x1, y1, x2, y2]} |
| 106 | keypoints_dict: {track_id: np.ndarray shape (17, 3)} where |
| 107 | columns are (x, y, confidence) |
| 108 | or None if the model did not return keypoints. |
| 109 | """ |
| 110 | results = self.model.track(frame, persist=True, verbose=False, conf=self.conf, iou=0.45, half=True, imgsz=self.imgsz)[0] |
| 111 | id_name_dict = results.names |
| 112 | |
| 113 | player_dict: dict[int, list] = {} |
| 114 | keypoints_dict: dict[int, np.ndarray | None] = {} |
| 115 | |
| 116 | if results.boxes is None or results.boxes.id is None: |
| 117 | return player_dict, keypoints_dict |
| 118 | |
| 119 | # Extract pose keypoints if available (pose model) |
| 120 | kps_data = None |
| 121 | if results.keypoints is not None and results.keypoints.data is not None: |
| 122 | kps_data = results.keypoints.data.cpu().numpy() # (N, 17, 3) |
| 123 | |
| 124 | for det_idx, box in enumerate(results.boxes): |
| 125 | track_id = int(box.id.tolist()[0]) |
| 126 | object_cls_name = id_name_dict[box.cls.tolist()[0]] |
| 127 | if object_cls_name != "person": |
| 128 | continue |
| 129 | |
| 130 | player_dict[track_id] = box.xyxy.tolist()[0] |
| 131 | |
| 132 | if kps_data is not None and det_idx < len(kps_data): |
no outgoing calls