|
|
@@ -13,6 +13,7 @@ import numpy as np
|
|
|
import threading
|
|
|
import queue
|
|
|
import time
|
|
|
+import logging
|
|
|
from typing import Optional, List, Tuple, Dict, Any
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@@ -20,6 +21,8 @@ from config import PANORAMA_CAMERA, DETECTION_CONFIG
|
|
|
from dahua_sdk import DahuaSDK, PTZCommand
|
|
|
from video_lock import safe_read, safe_is_opened
|
|
|
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
|
|
|
@dataclass
|
|
|
class DetectedObject:
|
|
|
@@ -582,33 +585,43 @@ class ObjectDetector:
|
|
|
results.append(obj)
|
|
|
|
|
|
except Exception as e:
|
|
|
- print(f"RKNN/ONNX 检测错误: {e}")
|
|
|
+ logger.error(f"RKNN/ONNX 检测错误: {e}")
|
|
|
|
|
|
return results
|
|
|
|
|
|
def detect(self, frame: np.ndarray) -> List[DetectedObject]:
|
|
|
- """
|
|
|
- 使用YOLO11检测物体
|
|
|
- Args:
|
|
|
- frame: 输入图像
|
|
|
- Returns:
|
|
|
- 检测结果列表
|
|
|
- """
|
|
|
+ """检测物体返回所有类别结果"""
|
|
|
if frame is None:
|
|
|
return []
|
|
|
|
|
|
- # 优先使用 RKNN/ONNX 模型
|
|
|
if hasattr(self, 'rknn') and self.rknn is not None:
|
|
|
- return self._detect_rknn(frame)
|
|
|
+ results = self._detect_rknn(frame)
|
|
|
+ if results:
|
|
|
+ self._log_detections("RKNN", results, frame)
|
|
|
+ return results
|
|
|
elif hasattr(self, 'session') and self.session is not None:
|
|
|
- return self._detect_rknn(frame)
|
|
|
- # 使用 YOLO 模型
|
|
|
+ results = self._detect_rknn(frame)
|
|
|
+ if results:
|
|
|
+ self._log_detections("ONNX", results, frame)
|
|
|
+ return results
|
|
|
elif self.model is not None:
|
|
|
- return self._detect_yolo(frame)
|
|
|
+ results = self._detect_yolo(frame)
|
|
|
+ if results:
|
|
|
+ self._log_detections("YOLO", results, frame)
|
|
|
+ return results
|
|
|
else:
|
|
|
- print("[错误] 没有可用的检测模型")
|
|
|
+ logger.error("[YOLO] 没有可用的检测模型")
|
|
|
return []
|
|
|
|
|
|
+ def _log_detections(self, model_type: str, results: List[DetectedObject], frame: np.ndarray):
|
|
|
+ if not results:
|
|
|
+ return
|
|
|
+ class_counts = {}
|
|
|
+ for r in results:
|
|
|
+ class_counts[r.class_name] = class_counts.get(r.class_name, 0) + 1
|
|
|
+ h, w = frame.shape[:2]
|
|
|
+ logger.info(f"[YOLO] {model_type}: {len(results)}个目标 {class_counts} (帧尺寸={w}x{h})")
|
|
|
+
|
|
|
def _detect_yolo(self, frame: np.ndarray) -> List[DetectedObject]:
|
|
|
"""使用 YOLO 模型检测"""
|
|
|
results = []
|
|
|
@@ -655,7 +668,7 @@ class ObjectDetector:
|
|
|
results.append(obj)
|
|
|
|
|
|
except Exception as e:
|
|
|
- print(f"YOLO11检测错误: {e}")
|
|
|
+ logger.error(f"YOLO11检测错误: {e}")
|
|
|
|
|
|
return results
|
|
|
|
|
|
@@ -670,15 +683,10 @@ class ObjectDetector:
|
|
|
return self.detect(frame)
|
|
|
|
|
|
def detect_persons(self, frame: np.ndarray) -> List[DetectedObject]:
|
|
|
- """
|
|
|
- 检测人体
|
|
|
- Args:
|
|
|
- frame: 输入图像
|
|
|
- Returns:
|
|
|
- 检测到的人体列表
|
|
|
- """
|
|
|
- results = self.detect(frame)
|
|
|
- return [obj for obj in results if obj.class_name == 'person']
|
|
|
+ """检测人体(支持中英文类别名)"""
|
|
|
+ all_detections = self.detect(frame)
|
|
|
+ person_classes = {'person', '人'}
|
|
|
+ return [obj for obj in all_detections if obj.class_name in person_classes]
|
|
|
|
|
|
def release(self):
|
|
|
"""释放模型资源"""
|