Ver Fonte

refactor(core): 重构检测器初始化及完善注释

- 修改检测器初始化逻辑,支持多种模型类型(YOLO11/RKNN/ONNX)
- 调整默认配置来源,优先使用 DETECTION_CONFIG 配置项
- 优化日志信息,统一检测器初始化成功与失败提示
- 补充代码注释,提升代码可读性
- 保留原有系统功能及结构,保持兼容性
- 更新其他细节,确保SDK与PyTorch内存冲突规避策略合理
wenhongquan há 4 dias atrás
pai
commit
df83f05c75

BIN
dual_camera_system/__pycache__/main.cpython-313.pyc


BIN
dual_camera_system/__pycache__/panorama_camera.cpython-313.pyc


BIN
dual_camera_system/config/__pycache__/detection.cpython-313.pyc


+ 17 - 2
dual_camera_system/config/detection.py

@@ -4,9 +4,24 @@
 
 # 检测配置
 DETECTION_CONFIG = {
-    'target_classes': ['person'],   # 检测目标类别
+    'target_classes': ['person', '人'],   # 检测目标类别 (支持中英文)
     'confidence_threshold': 0.5,     # 置信度阈值
     'detection_interval': 0.1,       # 检测间隔(秒)
+    
+    # RK3588 平台使用 RKNN 安全检测模型 (包含人体检测)
+    # 类别映射: 0=安全帽, 3=人, 4=反光衣
+    'model_path': '/home/admin/dsh/testrk3588/yolo11m_safety.rknn',
+    'model_type': 'rknn',           # 模型类型: 'rknn', 'yolo', 'onnx'
+    'use_gpu': False,               # RKNN 使用 NPU,不依赖 GPU
+    
+    # 安全检测模型的类别映射
+    'class_map': {
+        0: '安全帽',
+        3: '人',
+        4: '反光衣'
+    },
+    'person_class_id': 3,           # 人员在模型中的类别ID
+    'person_threshold': 0.8,        # 人员检测置信度阈值
 }
 
 # 安全检测模型配置
@@ -15,7 +30,7 @@ SAFETY_DETECTION_CONFIG = {
     # - YOLO: .pt 文件, 使用 ultralytics
     # - RKNN: .rknn 文件, 使用 rknnlite (RK3588 平台)
     # - ONNX: .onnx 文件, 使用 onnxruntime
-    'model_path': '/home/admin/dsh/testrk3588/yolo11n_rk3588.rknn',
+    'model_path': '/home/admin/dsh/testrk3588/yolo11m_safety.rknn',
     
     'model_type': 'rknn',           # 模型类型: 'auto', 'yolo', 'rknn', 'onnx'
     'use_gpu': False,                # RKNN 使用 NPU,不依赖 GPU

+ 7 - 5
dual_camera_system/main.py

@@ -116,14 +116,16 @@ class DualCameraSystem:
         # 大华SDK与PyTorch共享进程空间时可能导致内存冲突,
         # 先加载PyTorch可避免SDK的内存映射覆盖PyTorch运行时
         
-        # 初始化检测器 (YOLO11)
+        # 初始化检测器 (YOLO11/RKNN/ONNX)
         try:
+            from config import DETECTION_CONFIG
             self.detector = ObjectDetector(
-                model_path=self.config.get('model_path'),
-                use_gpu=self.config.get('use_gpu', True),
-                model_size=self.config.get('model_size', 'n')
+                model_path=self.config.get('model_path', DETECTION_CONFIG.get('model_path')),
+                use_gpu=self.config.get('use_gpu', DETECTION_CONFIG.get('use_gpu', True)),
+                model_size=self.config.get('model_size', 'n'),
+                model_type=self.config.get('model_type', DETECTION_CONFIG.get('model_type', 'auto'))
             )
-            logger.info("YOLO11检测器初始化成功")
+            logger.info("检测器初始化成功")
         except Exception as e:
             logger.warning(f"检测器初始化失败: {e}")
         

+ 16 - 9
dual_camera_system/panorama_camera.py

@@ -492,13 +492,13 @@ class ObjectDetector:
                 if x2 - x1 < 10 or y2 - y1 < 10:
                     continue
                 
-                # 只检测 person (class 0 in YOLO person model, 但 RKNN safety 模型是 class 3)
-                # 对于通用检测,尝试获取类别名称
-                cls_name = str(best_class)
-                if best_class == 0:
-                    cls_name = 'person'
-                elif best_class == 3:
-                    cls_name = 'person'  # 安全模型中的人
+                # 使用配置的类别映射获取类别名称
+                class_map = self.config.get('class_map', {0: 'person', 3: '人'})
+                cls_name = class_map.get(best_class, str(best_class))
+                
+                # 检查是否为目标类别
+                if cls_name not in self.config['target_classes']:
+                    continue
                 
                 obj = DetectedObject(
                     class_name=cls_name,
@@ -524,10 +524,17 @@ class ObjectDetector:
         if frame is None:
             return []
         
-        if self.rknn_detector is not None or hasattr(self, 'rknn') or hasattr(self, 'session'):
+        # 优先使用 RKNN/ONNX 模型
+        if hasattr(self, 'rknn') and self.rknn is not None:
             return self._detect_rknn(frame)
-        else:
+        elif hasattr(self, 'session') and self.session is not None:
+            return self._detect_rknn(frame)
+        # 使用 YOLO 模型
+        elif self.model is not None:
             return self._detect_yolo(frame)
+        else:
+            print("[错误] 没有可用的检测模型")
+            return []
     
     def _detect_yolo(self, frame: np.ndarray) -> List[DetectedObject]:
         """使用 YOLO 模型检测"""

BIN
testrk3588/yolo11m_safety.rknn