Upload engines/trackers/person_tracker.py with huggingface_hub

engines/trackers/person_tracker.py

@@ -15,6 +15,14 @@ try:
 except ImportError:
     logger.warning("deep-sort-realtime 未安装，人物追踪功能将不可用。安装方式: pip install deep-sort-realtime")
 
+# 尝试导入 PaddleOCR，不可用时优雅降级
+_paddleocr_available = False
+try:
+    from paddleocr import PaddleOCR
+    _paddleocr_available = True
+except ImportError:
+    logger.debug("paddleocr 未安装，球衣号码识别功能将不可用。安装方式: pip install paddleocr")
+
 
 def _is_gpu_available() -> bool:
     try:
@@ -40,7 +48,7 @@ class PersonTracker:
         self.sample_fps = sample_fps
         self.gpu = _is_gpu_available()
 
-        # 追踪结果：{track_id: {frames: [frame_indices], bbox_avg: [cx, cy], embeddings: [emb1, emb2, ...]}}
+        # 追踪结果：{track_id: {frames: [frame_indices], bbox_avg: [cx, cy], embeddings: [emb1, emb2, ...], jersey_numbers: [str, ...]}}
         self.tracks: Dict[int, Dict[str, Any]] = {}
 
         # 视频信息
@@ -86,6 +94,85 @@ class PersonTracker:
             logger.warning(f"X-CLIP 模型加载失败: {e}")
             return None, None
 
+    def _get_ocr_engine(self):
+        """获取 PaddleOCR 引擎（懒加载）"""
+        if not _paddleocr_available:
+            return None
+        try:
+            ocr = PaddleOCR(
+                use_angle_cls=False,
+                lang="en",
+                use_gpu=self.gpu,
+                show_log=False,
+                det_model_dir=None,  # 使用默认模型
+                rec_model_dir=None,
+            )
+            return ocr
+        except Exception as e:
+            logger.warning(f"PaddleOCR 初始化失败: {e}")
+            return None
+
+    def _recognize_jersey_number(self, frame: np.ndarray, bbox: List[float], ocr_engine) -> Optional[str]:
+        """识别人物球衣号码
+
+        从人物上半身区域裁剪，用 PaddleOCR 识别数字。
+
+        Args:
+            frame: 原始帧 (H, W, 3) BGR
+            bbox: [left, top, width, height]
+            ocr_engine: PaddleOCR 实例
+
+        Returns:
+            识别到的号码字符串 或 None
+        """
+        if ocr_engine is None:
+            return None
+
+        try:
+            h, w = frame.shape[:2]
+            left = max(0, int(bbox[0]))
+            top = max(0, int(bbox[1]))
+            right = min(w, int(bbox[0] + bbox[2]))
+            bottom = min(h, int(bbox[1] + bbox[3]))
+
+            if right - left < 20 or bottom - top < 40:
+                return None
+
+            # 裁剪上半身区域（球衣号码通常在胸部位置）
+            torso_top = top
+            torso_bottom = top + int((bottom - top) * 0.6)  # 上 60%
+            torso_left = left + int((right - left) * 0.15)  # 去掉边缘
+            torso_right = right - int((right - left) * 0.15)
+
+            if torso_right - torso_left < 10 or torso_bottom - torso_top < 10:
+                return None
+
+            crop = frame[torso_top:torso_bottom, torso_left:torso_right]
+            if crop.size == 0:
+                return None
+
+            # PaddleOCR 识别
+            result = ocr_engine.ocr(crop, cls=False)
+            if not result or not result[0]:
+                return None
+
+            # 从识别结果中提取纯数字文本
+            for line in result[0]:
+                text = line[1][0].strip()
+                confidence = line[1][1]
+                # 过滤：只接受 1-2 位纯数字，置信度 > 0.5
+                if confidence < 0.5:
+                    continue
+                digits = "".join(c for c in text if c.isdigit())
+                if 1 <= len(digits) <= 2 and digits != "0":
+                    return digits
+
+            return None
+
+        except Exception as e:
+            logger.debug(f"球衣号码识别失败: {e}")
+            return None
+
     def _encode_person_crop(self, frame: np.ndarray, bbox: List[float],
                            xclip_model, xclip_processor) -> Optional[np.ndarray]:
         """裁剪人物区域并用 X-CLIP 视觉编码器提取嵌入
@@ -186,6 +273,14 @@ class PersonTracker:
         # 初始化 DeepSort 追踪器
         tracker = DeepSort(max_age=30, n_init=3, nn_budget=100)
 
+        # 初始化 PaddleOCR（球衣号码识别）
+        ocr_engine = None
+        if _paddleocr_available:
+            _report(0.12, "初始化 PaddleOCR 球衣号码识别...")
+            ocr_engine = self._get_ocr_engine()
+            if ocr_engine is None:
+                logger.info("PaddleOCR 不可用，跳过球衣号码识别")
+
         cap = cv2.VideoCapture(self.video_path)
         if not cap.isOpened():
             logger.error("无法打开视频文件")
@@ -243,6 +338,7 @@ class PersonTracker:
                                 "frames": [],
                                 "bbox_avg": [0.0, 0.0],
                                 "embeddings": [],
+                                "jersey_numbers": [],
                                 "_cx_sum": 0.0,
                                 "_cy_sum": 0.0,
                             }
@@ -265,6 +361,14 @@ class PersonTracker:
                             if embedding is not None:
                                 self.tracks[track_id]["embeddings"].append(embedding)
 
+                        # 球衣号码识别（每 10 帧采样一次，OCR 较慢）
+                        if ocr_engine is not None and n % 10 == 1:
+                            jersey_num = self._recognize_jersey_number(
+                                frame, [ltrb[0], ltrb[1], bbox_w, bbox_h], ocr_engine,
+                            )
+                            if jersey_num is not None:
+                                self.tracks[track_id]["jersey_numbers"].append(jersey_num)
+
                 except Exception as e:
                     logger.warning(f"DeepSort 追踪失败 (frame {frame_idx}): {e}")
 
@@ -284,10 +388,18 @@ class PersonTracker:
 
         cap.release()
 
-        # 清理临时字段
+        # 清理临时字段，去重球衣号码
         for track_id in self.tracks:
             self.tracks[track_id].pop("_cx_sum", None)
             self.tracks[track_id].pop("_cy_sum", None)
+            # 球衣号码去重，保留出现次数最多的
+            jns = self.tracks[track_id].get("jersey_numbers", [])
+            if jns:
+                from collections import Counter
+                most_common = Counter(jns).most_common(1)[0][0]
+                self.tracks[track_id]["jersey_number"] = most_common
+            else:
+                self.tracks[track_id]["jersey_number"] = None
 
         # 过滤掉出现次数太少的追踪（少于 5 帧）
         self.tracks = {