Parcourir la source

feat(coordinator): 添加联动控制器和事件驱动协调器功能

refactor(detection): 重构检测模型支持RKNN/ONNX/YOLO多种格式

fix(camera): 修复RTSP连接配置和SDK路径问题

style(system): 更新系统配置默认值

perf(panorama): 优化人体分割和OCR预处理流程

docs(config): 添加配置模块文档字符串

test(ocr): 增加本地OCR识别器作为备用方案

chore(dahua_sdk): 更新大华SDK封装和错误处理
wenhongquan il y a 5 jours
Parent
commit
2a19b050da

BIN
dual_camera_system/__pycache__/coordinator.cpython-310.pyc


BIN
dual_camera_system/__pycache__/dahua_sdk.cpython-310.pyc


BIN
dual_camera_system/__pycache__/ocr_recognizer.cpython-310.pyc


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


BIN
dual_camera_system/__pycache__/ptz_camera.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/__init__.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/camera.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/coordinator.cpython-310.pyc


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


BIN
dual_camera_system/config/__pycache__/event.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/llm.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/ocr.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/ptz.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/system.cpython-310.pyc


BIN
dual_camera_system/config/__pycache__/voice.cpython-310.pyc


+ 11 - 9
dual_camera_system/config/camera.py

@@ -13,24 +13,26 @@ LOG_CONFIG = {
 
 
 # 全景摄像头配置
 # 全景摄像头配置
 PANORAMA_CAMERA = {
 PANORAMA_CAMERA = {
-    'ip': '192.168.1.100',      # 全景摄像头IP
-    'port': 37777,              # 大华SDK默认端口
+    'ip': '192.168.20.196',      # 全景摄像头IP
+    'port': 554,                  # RTSP端口
     'username': 'admin',
     'username': 'admin',
-    'password': 'admin123',
-    'channel': 0,               # 通道号
+    'password': 'Aa1234567',
+    'channel': 1,                # 通道号
+    'rtsp_url': 'rtsp://admin:Aa1234567@192.168.20.196:554/cam/realmonitor?channel=1&subtype=1',
 }
 }
 
 
 # 球机配置
 # 球机配置
 PTZ_CAMERA = {
 PTZ_CAMERA = {
-    'ip': '192.168.1.101',      # 球机IP
-    'port': 37777,
+    'ip': '192.168.20.197',      # 球机IP
+    'port': 554,                 # RTSP端口
     'username': 'admin',
     'username': 'admin',
-    'password': 'admin123',
-    'channel': 0,
+    'password': 'Aa1234567',
+    'channel': 1,
+    'rtsp_url': 'rtsp://admin:Aa1234567@192.168.20.197:554/cam/realmonitor?channel=1&subtype=1',
 }
 }
 
 
 # 大华SDK库路径
 # 大华SDK库路径
 SDK_PATH = {
 SDK_PATH = {
-    'lib_path': '/home/wen/dsh/dh/Bin',
+    'lib_path': '/Users/wenhongquan/Desktop/阿里云同步/项目/dnn/德胜河 AI/dsh/dh/Bin',
     'netsdk': 'libdhnetsdk.so',
     'netsdk': 'libdhnetsdk.so',
 }
 }

+ 8 - 2
dual_camera_system/config/detection.py

@@ -11,8 +11,14 @@ DETECTION_CONFIG = {
 
 
 # 安全检测模型配置
 # 安全检测模型配置
 SAFETY_DETECTION_CONFIG = {
 SAFETY_DETECTION_CONFIG = {
-    'model_path': '/home/wen/dsh/yolo/yolo11m_safety.pt',  # 安全检测模型路径
-    'use_gpu': True,                 # 是否使用 GPU
+    # 模型路径 - 支持三种格式:
+    # - YOLO: .pt 文件, 使用 ultralytics
+    # - RKNN: .rknn 文件, 使用 rknnlite (RK3588 平台)
+    # - ONNX: .onnx 文件, 使用 onnxruntime
+    'model_path': '/Users/wenhongquan/Desktop/阿里云同步/项目/dnn/德胜河 AI/dsh/testrk3588/yolo11n_rk3588.pt',
+    
+    'model_type': 'yolo',           # 模型类型: 'auto', 'yolo', 'rknn', 'onnx'
+    'use_gpu': False,                # RKNN 使用 NPU,不依赖 GPU
     'conf_threshold': 0.5,           # 一般物品置信度阈值 (安全帽、反光衣)
     'conf_threshold': 0.5,           # 一般物品置信度阈值 (安全帽、反光衣)
     'person_threshold': 0.8,         # 人员检测置信度阈值
     'person_threshold': 0.8,         # 人员检测置信度阈值
     
     

+ 4 - 4
dual_camera_system/config/system.py

@@ -24,12 +24,12 @@ SYSTEM_CONFIG = {
     'enable_ptz_tracking': True,         # 启用 PTZ 跟踪联动
     'enable_ptz_tracking': True,         # 启用 PTZ 跟踪联动
     
     
     # OCR 与大模型
     # OCR 与大模型
-    'enable_ocr': True,                  # 启用 OCR 编号识别
-    'enable_llm': True,                  # 启用大模型判断
+    'enable_ocr': False,                  # 启用 OCR 编号识别
+    'enable_llm': False,                  # 启用大模型判断
     
     
     # 事件与播报
     # 事件与播报
-    'enable_event_push': True,           # 启用事件推送
-    'enable_voice_announce': True,       # 启用语音播报
+    'enable_event_push': False,           # 启用事件推送
+    'enable_voice_announce': False,       # 启用语音播报
     
     
     # === 工作模式 ===
     # === 工作模式 ===
     'mode': 'safety',                    # 工作模式: 'safety'(安全检测), 'ocr'(编号识别)
     'mode': 'safety',                    # 工作模式: 'safety'(安全检测), 'ocr'(编号识别)

+ 166 - 19
dual_camera_system/panorama_camera.py

@@ -8,6 +8,7 @@ import numpy as np
 import threading
 import threading
 import queue
 import queue
 import time
 import time
+import os
 from typing import Optional, List, Tuple, Dict, Any
 from typing import Optional, List, Tuple, Dict, Any
 from dataclasses import dataclass
 from dataclasses import dataclass
 
 
@@ -330,41 +331,95 @@ class ObjectDetector:
     """
     """
     物体检测器
     物体检测器
     使用YOLO11模型进行人体检测
     使用YOLO11模型进行人体检测
+    支持 YOLO (.pt), RKNN (.rknn), ONNX (.onnx) 模型
     """
     """
     
     
-    def __init__(self, model_path: str = None, use_gpu: bool = True, model_size: str = 'n'):
+    def __init__(self, model_path: str = None, use_gpu: bool = True, model_size: str = 'n',
+                 model_type: str = 'auto'):
         """
         """
         初始化检测器
         初始化检测器
         Args:
         Args:
-            model_path: 模型路径 (自定义模型)
+            model_path: 模型路径 (支持 .pt, .rknn, .onnx)
             use_gpu: 是否使用GPU
             use_gpu: 是否使用GPU
-            model_size: 模型尺寸 ('n', 's', 'm', 'l', 'x')
+            model_size: 模型尺寸 ('n', 's', 'm', 'l', 'x') - 仅 YOLO 模型有效
+            model_type: 模型类型 ('auto', 'yolo', 'rknn', 'onnx')
         """
         """
         self.model = None
         self.model = None
+        self.rknn_detector = None
         self.model_path = model_path
         self.model_path = model_path
         self.use_gpu = use_gpu
         self.use_gpu = use_gpu
         self.model_size = model_size
         self.model_size = model_size
+        self.model_type = model_type
         self.config = DETECTION_CONFIG
         self.config = DETECTION_CONFIG
         self.device = 'cuda:0' if use_gpu else 'cpu'
         self.device = 'cuda:0' if use_gpu else 'cpu'
         
         
+        # 根据扩展名自动判断模型类型
+        if model_path:
+            ext = os.path.splitext(model_path)[1].lower()
+            if ext == '.rknn':
+                self.model_type = 'rknn'
+            elif ext == '.onnx':
+                self.model_type = 'onnx'
+            elif ext == '.pt':
+                self.model_type = 'yolo'
+        
         self._load_model()
         self._load_model()
     
     
     def _load_model(self):
     def _load_model(self):
+        """加载检测模型"""
+        if self.model_type == 'rknn':
+            self._load_rknn_model()
+        elif self.model_type == 'onnx':
+            self._load_onnx_model()
+        else:
+            self._load_yolo_model()
+    
+    def _load_rknn_model(self):
+        """加载 RKNN 模型"""
+        if not self.model_path:
+            raise ValueError("RKNN 模型需要指定 model_path")
+        
+        try:
+            from rknnlite.api import RKNNLite
+            
+            self.rknn = RKNNLite()
+            ret = self.rknn.load_rknn(self.model_path)
+            if ret != 0:
+                raise RuntimeError(f"加载 RKNN 模型失败: {self.model_path}")
+            
+            ret = self.rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
+            if ret != 0:
+                raise RuntimeError(f"初始化 RKNN 运行时失败")
+            
+            print(f"RKNN 模型加载成功: {self.model_path}")
+        except ImportError:
+            raise ImportError("未安装 rknnlite,请运行: pip install rknnlite2")
+    
+    def _load_onnx_model(self):
+        """加载 ONNX 模型"""
+        if not self.model_path:
+            raise ValueError("ONNX 模型需要指定 model_path")
+        
+        try:
+            import onnxruntime as ort
+            self.session = ort.InferenceSession(self.model_path)
+            self.input_name = self.session.get_inputs()[0].name
+            self.output_name = self.session.get_outputs()[0].name
+            print(f"ONNX 模型加载成功: {self.model_path}")
+        except ImportError:
+            raise ImportError("未安装 onnxruntime,请运行: pip install onnxruntime")
+    
+    def _load_yolo_model(self):
         """加载YOLO11检测模型"""
         """加载YOLO11检测模型"""
         try:
         try:
-            # 使用ultralytics YOLO11
             from ultralytics import YOLO
             from ultralytics import YOLO
             
             
             if self.model_path:
             if self.model_path:
-                # 使用自定义模型
                 self.model = YOLO(self.model_path)
                 self.model = YOLO(self.model_path)
             else:
             else:
-                # 使用YOLO11预训练模型
-                # YOLO11模型命名: yolo11n.pt, yolo11s.pt, yolo11m.pt, yolo11l.pt, yolo11x.pt
                 model_name = f'yolo11{self.model_size}.pt'
                 model_name = f'yolo11{self.model_size}.pt'
                 self.model = YOLO(model_name)
                 self.model = YOLO(model_name)
             
             
-            # 预热模型
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             self.model(dummy, device=self.device, verbose=False)
             self.model(dummy, device=self.device, verbose=False)
             
             
@@ -378,9 +433,95 @@ class ObjectDetector:
     
     
     def _load_opencv_model(self):
     def _load_opencv_model(self):
         """使用OpenCV加载模型"""
         """使用OpenCV加载模型"""
-        # 可以加载ONNX模型
         pass
         pass
     
     
+    def _letterbox(self, image, size=(640, 640)):
+        """Letterbox 预处理"""
+        h0, w0 = image.shape[:2]
+        ih, iw = size
+        scale = min(iw / w0, ih / h0)
+        new_w, new_h = int(w0 * scale), int(h0 * scale)
+        pad_w = (iw - new_w) // 2
+        pad_h = (ih - new_h) // 2
+        resized = cv2.resize(image, (new_w, new_h))
+        canvas = np.full((ih, iw, 3), 114, dtype=np.uint8)
+        canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
+        return canvas, scale, pad_w, pad_h, h0, w0
+    
+    def _detect_rknn(self, frame: np.ndarray) -> List[DetectedObject]:
+        """使用 RKNN/ONNX 模型检测"""
+        results = []
+        
+        try:
+            canvas, scale, pad_w, pad_h, h0, w0 = self._letterbox(frame)
+            
+            if hasattr(self, 'rknn'):
+                # RKNN
+                img = canvas[..., ::-1].astype(np.float32) / 255.0
+                blob = img[None, ...]
+                outputs = self.rknn.inference(inputs=[blob])
+            else:
+                # ONNX
+                img = canvas[..., ::-1].astype(np.float32) / 255.0
+                img = img.transpose(2, 0, 1)
+                blob = img[None, ...]
+                outputs = self.session.run([self.output_name], {self.input_name: blob})
+            
+            output = outputs[0]
+            if len(output.shape) == 3:
+                output = output[0]
+            
+            num_boxes = output.shape[1]
+            conf_threshold = self.config['confidence_threshold']
+            
+            for i in range(num_boxes):
+                x_center = float(output[0, i])
+                y_center = float(output[1, i])
+                width = float(output[2, i])
+                height = float(output[3, i])
+                
+                class_probs = output[4:, i]
+                best_class = int(np.argmax(class_probs))
+                confidence = float(class_probs[best_class])
+                
+                if confidence < conf_threshold:
+                    continue
+                
+                # 转换到原始图像坐标
+                x1 = int(((x_center - width / 2) - pad_w) / scale)
+                y1 = int(((y_center - height / 2) - pad_h) / scale)
+                x2 = int(((x_center + width / 2) - pad_w) / scale)
+                y2 = int(((y_center + height / 2) - pad_h) / scale)
+                
+                x1 = max(0, min(w0, x1))
+                y1 = max(0, min(h0, y1))
+                x2 = max(0, min(w0, x2))
+                y2 = max(0, min(h0, y2))
+                
+                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'  # 安全模型中的人
+                
+                obj = DetectedObject(
+                    class_name=cls_name,
+                    confidence=confidence,
+                    bbox=(x1, y1, x2 - x1, y2 - y1),
+                    center=((x1 + x2) // 2, (y1 + y2) // 2)
+                )
+                results.append(obj)
+                
+        except Exception as e:
+            print(f"RKNN/ONNX 检测错误: {e}")
+        
+        return results
+    
     def detect(self, frame: np.ndarray) -> List[DetectedObject]:
     def detect(self, frame: np.ndarray) -> List[DetectedObject]:
         """
         """
         使用YOLO11检测物体
         使用YOLO11检测物体
@@ -389,13 +530,19 @@ class ObjectDetector:
         Returns:
         Returns:
             检测结果列表
             检测结果列表
         """
         """
-        if self.model is None or frame is None:
+        if frame is None:
             return []
             return []
         
         
+        if self.rknn_detector is not None or hasattr(self, 'rknn') or hasattr(self, 'session'):
+            return self._detect_rknn(frame)
+        else:
+            return self._detect_yolo(frame)
+    
+    def _detect_yolo(self, frame: np.ndarray) -> List[DetectedObject]:
+        """使用 YOLO 模型检测"""
         results = []
         results = []
         
         
         try:
         try:
-            # YOLO11推理
             detections = self.model(
             detections = self.model(
                 frame, 
                 frame, 
                 device=self.device, 
                 device=self.device, 
@@ -409,28 +556,22 @@ class ObjectDetector:
                     continue
                     continue
                     
                     
                 for i in range(len(boxes)):
                 for i in range(len(boxes)):
-                    # 获取类别
                     cls_id = int(boxes.cls[i])
                     cls_id = int(boxes.cls[i])
                     cls_name = det.names[cls_id]
                     cls_name = det.names[cls_id]
                     
                     
-                    # 过滤目标类别
                     if cls_name not in self.config['target_classes']:
                     if cls_name not in self.config['target_classes']:
                         continue
                         continue
                     
                     
-                    # 获取置信度
                     conf = float(boxes.conf[i])
                     conf = float(boxes.conf[i])
                     
                     
-                    # 获取边界框
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     x1, y1, x2, y2 = map(int, xyxy)
                     x1, y1, x2, y2 = map(int, xyxy)
                     width = x2 - x1
                     width = x2 - x1
                     height = y2 - y1
                     height = y2 - y1
                     
                     
-                    # 过滤过小的检测框
                     if width < 10 or height < 10:
                     if width < 10 or height < 10:
                         continue
                         continue
                     
                     
-                    # 计算中心点
                     center_x = x1 + width // 2
                     center_x = x1 + width // 2
                     center_y = y1 + height // 2
                     center_y = y1 + height // 2
                     
                     
@@ -455,8 +596,6 @@ class ObjectDetector:
         Returns:
         Returns:
             带关键点的检测结果列表
             带关键点的检测结果列表
         """
         """
-        # 如果使用pose模型,可以获取人体关键点
-        # 用于更精确的人体定位
         return self.detect(frame)
         return self.detect(frame)
     
     
     def detect_persons(self, frame: np.ndarray) -> List[DetectedObject]:
     def detect_persons(self, frame: np.ndarray) -> List[DetectedObject]:
@@ -469,6 +608,14 @@ class ObjectDetector:
         """
         """
         results = self.detect(frame)
         results = self.detect(frame)
         return [obj for obj in results if obj.class_name == 'person']
         return [obj for obj in results if obj.class_name == 'person']
+    
+    def release(self):
+        """释放模型资源"""
+        if hasattr(self, 'rknn') and self.rknn:
+            self.rknn.release()
+            self.rknn = None
+        self.model = None
+        self.session = None
 
 
 
 
 class PersonTracker:
 class PersonTracker:

+ 352 - 25
dual_camera_system/safety_detector.py

@@ -2,6 +2,10 @@
 施工现场安全行为检测模块
 施工现场安全行为检测模块
 使用 YOLO11 模型检测人员、安全帽、反光衣
 使用 YOLO11 模型检测人员、安全帽、反光衣
 判断是否存在违规行为(未戴安全帽、未穿反光衣)
 判断是否存在违规行为(未戴安全帽、未穿反光衣)
+
+支持两种模型格式:
+- YOLO (.pt/.onnx): 使用 ultralytics 库
+- RKNN (.rknn): 使用 rknnlite 库 (RK3588 平台)
 """
 """
 
 
 import cv2
 import cv2
@@ -9,6 +13,239 @@ import numpy as np
 from typing import Optional, List, Tuple, Dict, Any
 from typing import Optional, List, Tuple, Dict, Any
 from dataclasses import dataclass
 from dataclasses import dataclass
 from enum import Enum
 from enum import Enum
+import os
+
+
+# ============================================
+# RKNN 模型支持
+# ============================================
+
+@dataclass
+class Detection:
+    """检测结果 (用于 RKNN 模型)"""
+    class_id: int
+    class_name: str
+    confidence: float
+    bbox: Tuple[int, int, int, int]
+
+
+def nms(dets, iou_threshold=0.45):
+    """非极大值抑制"""
+    if len(dets) == 0:
+        return []
+    
+    boxes = np.array([[d.bbox[0], d.bbox[1], d.bbox[2], d.bbox[3], d.confidence] for d in dets])
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    scores = boxes[:, 4]
+    
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+    
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+        
+        inds = np.where(ovr <= iou_threshold)[0]
+        order = order[inds + 1]
+    
+    return [dets[i] for i in keep]
+
+
+class BaseDetector:
+    """检测器基类 (用于 RKNN/ONNX 模型)"""
+    
+    # 类别映射: 0: 安全帽, 3: 人, 4: 反光衣
+    LABEL_MAP = {0: '安全帽', 4: '安全衣', 3: '人'}
+    
+    def __init__(self):
+        self.input_size = (640, 640)
+        self.num_classes = 5
+    
+    def letterbox(self, image):
+        """Letterbox 预处理,保持宽高比"""
+        h0, w0 = image.shape[:2]
+        ih, iw = self.input_size
+        scale = min(iw / w0, ih / h0)
+        new_w, new_h = int(w0 * scale), int(h0 * scale)
+        pad_w = (iw - new_w) // 2
+        pad_h = (ih - new_h) // 2
+        resized = cv2.resize(image, (new_w, new_h))
+        canvas = np.full((ih, iw, 3), 114, dtype=np.uint8)
+        canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
+        return canvas, scale, pad_w, pad_h, h0, w0
+    
+    def postprocess(self, outputs, scale, pad_w, pad_h, h0, w0, conf_threshold_map):
+        """后处理"""
+        dets = []
+        
+        if not outputs:
+            return dets
+        
+        output = outputs[0]
+        
+        if len(output.shape) == 3:
+            output = output[0]
+        
+        num_boxes = output.shape[1]
+        
+        for i in range(num_boxes):
+            x_center = float(output[0, i])
+            y_center = float(output[1, i])
+            width = float(output[2, i])
+            height = float(output[3, i])
+            
+            class_probs = output[4:4+self.num_classes, i]
+            best_class = int(np.argmax(class_probs))
+            confidence = float(class_probs[best_class])
+            
+            if best_class not in self.LABEL_MAP:
+                continue
+            
+            conf_threshold = conf_threshold_map.get(best_class, 0.5)
+            
+            if confidence < conf_threshold:
+                continue
+            
+            # 移除 padding 并缩放到原始图像尺寸
+            x1 = int(((x_center - width / 2) - pad_w) / scale)
+            y1 = int(((y_center - height / 2) - pad_h) / scale)
+            x2 = int(((x_center + width / 2) - pad_w) / scale)
+            y2 = int(((y_center + height / 2) - pad_h) / scale)
+            
+            x1 = max(0, min(w0, x1))
+            y1 = max(0, min(h0, y1))
+            x2 = max(0, min(w0, x2))
+            y2 = max(0, min(h0, y2))
+            
+            det = Detection(
+                class_id=best_class,
+                class_name=self.LABEL_MAP[best_class],
+                confidence=confidence,
+                bbox=(x1, y1, x2, y2)
+            )
+            dets.append(det)
+        
+        dets = nms(dets, iou_threshold=0.45)
+        return dets
+    
+    def detect(self, image, conf_threshold_map):
+        raise NotImplementedError
+    
+    def release(self):
+        pass
+
+
+class RKNNDetector(BaseDetector):
+    """RKNN 检测器 - 使用 NHWC 输入格式 (1, H, W, C)"""
+    
+    def __init__(self, model_path: str):
+        super().__init__()
+        self.model_path = model_path
+        self.rknn = None
+        
+        try:
+            from rknnlite.api import RKNNLite
+            self.rknn = RKNNLite()
+        except ImportError:
+            raise ImportError("未安装 rknnlite,请运行: pip install rknnlite2 或参考 testrk3588/setup_rknn.sh")
+        
+        ret = self.rknn.load_rknn(model_path)
+        if ret != 0:
+            raise RuntimeError(f"加载 RKNN 模型失败: {model_path}")
+        
+        ret = self.rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
+        if ret != 0:
+            raise RuntimeError(f"初始化 RKNN 运行时失败")
+        
+        print(f"RKNN 模型加载成功: {model_path}")
+    
+    def detect(self, image, conf_threshold_map):
+        canvas, scale, pad_w, pad_h, h0, w0 = self.letterbox(image)
+        # RKNN 期望 NHWC (1, H, W, C), RGB, 归一化 0-1
+        img = canvas[..., ::-1].astype(np.float32) / 255.0
+        blob = img[None, ...]  # (1, 640, 640, 3)
+        outs = self.rknn.inference(inputs=[blob])
+        return self.postprocess(outs, scale, pad_w, pad_h, h0, w0, conf_threshold_map)
+    
+    def release(self):
+        if self.rknn:
+            self.rknn.release()
+            self.rknn = None
+
+
+class ONNXDetector(BaseDetector):
+    """ONNX 检测器 - 使用 NCHW 输入格式 (1, C, H, W)"""
+    
+    def __init__(self, model_path: str):
+        super().__init__()
+        self.model_path = model_path
+        
+        try:
+            import onnxruntime as ort
+            self.session = ort.InferenceSession(model_path)
+            self.input_name = self.session.get_inputs()[0].name
+            self.output_name = self.session.get_outputs()[0].name
+            print(f"ONNX 模型加载成功: {model_path}")
+        except ImportError:
+            raise ImportError("未安装 onnxruntime,请运行: pip install onnxruntime")
+        except Exception as e:
+            raise RuntimeError(f"加载 ONNX 模型失败: {e}")
+    
+    def detect(self, image, conf_threshold_map):
+        canvas, scale, pad_w, pad_h, h0, w0 = self.letterbox(image)
+        # ONNX 期望 NCHW (1, C, H, W), RGB, 归一化 0-1
+        img = canvas[..., ::-1].astype(np.float32) / 255.0
+        img = img.transpose(2, 0, 1)
+        blob = img[None, ...]  # (1, 3, 640, 640)
+        outs = self.session.run([self.output_name], {self.input_name: blob})
+        return self.postprocess(outs, scale, pad_w, pad_h, h0, w0, conf_threshold_map)
+    
+    def release(self):
+        self.session = None
+
+
+def create_detector(model_path: str):
+    """
+    创建检测器工厂函数
+    
+    Args:
+        model_path: 模型路径 (.rknn, .onnx, .pt)
+        
+    Returns:
+        检测器实例
+    """
+    ext = os.path.splitext(model_path)[1].lower()
+    
+    if ext == '.rknn':
+        print(f"使用 RKNN 模型: {model_path}")
+        return RKNNDetector(model_path)
+    elif ext == '.onnx':
+        print(f"使用 ONNX 模型: {model_path}")
+        return ONNXDetector(model_path)
+    elif ext == '.pt':
+        print(f"使用 YOLO 模型: {model_path}")
+        return None  # YOLO 使用原来的 SafetyDetector
+    else:
+        raise ValueError(f"不支持的模型格式: {ext}")
+
+
+# ============================================
+# 原有 YOLO 安全检测器
+# ============================================
 
 
 
 
 class SafetyViolationType(Enum):
 class SafetyViolationType(Enum):
@@ -81,15 +318,12 @@ class SafetyDetector:
     使用 YOLO11 检测人员、安全帽、反光衣
     使用 YOLO11 检测人员、安全帽、反光衣
     """
     """
     
     
-    # 类别映射 (根据 yolo11m_safety.pt 模型的训练标签)
-    # 0: 安全帽, 3: 人, 4: 安全衣/反光衣
     CLASS_MAP = {
     CLASS_MAP = {
         0: '安全帽',
         0: '安全帽',
         3: '人',
         3: '人',
         4: '反光衣'
         4: '反光衣'
     }
     }
     
     
-    # 反向映射
     CLASS_ID_MAP = {
     CLASS_ID_MAP = {
         'helmet': 0,
         'helmet': 0,
         'person': 3,
         'person': 3,
@@ -97,44 +331,94 @@ class SafetyDetector:
     }
     }
     
     
     def __init__(self, model_path: str = None, use_gpu: bool = True, 
     def __init__(self, model_path: str = None, use_gpu: bool = True, 
-                 conf_threshold: float = 0.5, person_threshold: float = 0.8):
+                 conf_threshold: float = 0.5, person_threshold: float = 0.8,
+                 model_type: str = 'auto'):
         """
         """
         初始化安全检测器
         初始化安全检测器
         
         
         Args:
         Args:
-            model_path: 模型路径,默认使用 yolo11m_safety.pt
-            use_gpu: 是否使用 GPU
+            model_path: 模型路径,默认使用 yolo11m_safety.pt 或 .rknn
+            use_gpu: 是否使用 GPU (仅 YOLO 模型有效)
             conf_threshold: 一般物品置信度阈值 (安全帽、反光衣)
             conf_threshold: 一般物品置信度阈值 (安全帽、反光衣)
             person_threshold: 人员检测置信度阈值
             person_threshold: 人员检测置信度阈值
+            model_type: 模型类型 ('auto', 'yolo', 'rknn', 'onnx')
         """
         """
         self.model = None
         self.model = None
-        self.model_path = model_path or '/home/wen/dsh/yolo/yolo11m_safety.pt'
+        self.rknn_detector = None
+        self.model_type = model_type
+        
+        # 根据扩展名自动判断模型类型
+        if model_path:
+            ext = os.path.splitext(model_path)[1].lower()
+            if ext == '.rknn':
+                self.model_type = 'rknn'
+            elif ext == '.onnx':
+                self.model_type = 'onnx'
+            elif ext == '.pt':
+                self.model_type = 'yolo'
+        
+        self.model_path = model_path
         self.use_gpu = use_gpu
         self.use_gpu = use_gpu
         self.device = 'cuda:0' if use_gpu else 'cpu'
         self.device = 'cuda:0' if use_gpu else 'cpu'
         
         
-        # 置信度阈值
         self.conf_threshold = conf_threshold
         self.conf_threshold = conf_threshold
         self.person_threshold = person_threshold
         self.person_threshold = person_threshold
         
         
-        # 加载模型
         self._load_model()
         self._load_model()
     
     
     def _load_model(self):
     def _load_model(self):
+        """加载检测模型"""
+        if self.model_type == 'rknn':
+            self._load_rknn_model()
+        elif self.model_type == 'onnx':
+            self._load_onnx_model()
+        else:
+            self._load_yolo_model()
+    
+    def _load_rknn_model(self):
+        """加载 RKNN 模型"""
+        if not self.model_path:
+            raise ValueError("RKNN 模型需要指定 model_path")
+        
+        try:
+            self.rknn_detector = RKNNDetector(self.model_path)
+            print(f"RKNN 安全检测模型加载成功: {self.model_path}")
+        except ImportError as e:
+            raise ImportError(f"rknnlite 未安装: {e}")
+        except Exception as e:
+            raise RuntimeError(f"加载 RKNN 模型失败: {e}")
+    
+    def _load_onnx_model(self):
+        """加载 ONNX 模型"""
+        if not self.model_path:
+            raise ValueError("ONNX 模型需要指定 model_path")
+        
+        try:
+            self.rknn_detector = ONNXDetector(self.model_path)
+            print(f"ONNX 安全检测模型加载成功: {self.model_path}")
+        except ImportError as e:
+            raise ImportError(f"onnxruntime 未安装: {e}")
+        except Exception as e:
+            raise RuntimeError(f"加载 ONNX 模型失败: {e}")
+    
+    def _load_yolo_model(self):
         """加载 YOLO11 安全检测模型"""
         """加载 YOLO11 安全检测模型"""
         try:
         try:
             from ultralytics import YOLO
             from ultralytics import YOLO
             
             
+            if not self.model_path:
+                self.model_path = '/home/wen/dsh/yolo/yolo11m_safety.pt'
+            
             self.model = YOLO(self.model_path)
             self.model = YOLO(self.model_path)
             
             
-            # 预热模型
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             self.model(dummy, device=self.device, verbose=False)
             self.model(dummy, device=self.device, verbose=False)
             
             
-            print(f"安全检测模型加载成功: {self.model_path} (device={self.device})")
+            print(f"YOLO 安全检测模型加载成功: {self.model_path} (device={self.device})")
         except ImportError:
         except ImportError:
             raise ImportError("未安装 ultralytics,请运行: pip install ultralytics")
             raise ImportError("未安装 ultralytics,请运行: pip install ultralytics")
         except Exception as e:
         except Exception as e:
-            raise RuntimeError(f"加载模型失败: {e}")
+            raise RuntimeError(f"加载 YOLO 模型失败: {e}")
     
     
     def detect(self, frame: np.ndarray) -> List[SafetyDetection]:
     def detect(self, frame: np.ndarray) -> List[SafetyDetection]:
         """
         """
@@ -146,9 +430,48 @@ class SafetyDetector:
         Returns:
         Returns:
             检测结果列表
             检测结果列表
         """
         """
-        if self.model is None or frame is None:
+        if frame is None:
             return []
             return []
         
         
+        if self.rknn_detector is not None:
+            return self._detect_rknn(frame)
+        else:
+            return self._detect_yolo(frame)
+    
+    def _detect_rknn(self, frame: np.ndarray) -> List[SafetyDetection]:
+        """使用 RKNN/ONNX 模型检测"""
+        results = []
+        
+        try:
+            conf_threshold_map = {
+                3: self.person_threshold,
+                0: self.conf_threshold,
+                4: self.conf_threshold
+            }
+            
+            detections = self.rknn_detector.detect(frame, conf_threshold_map)
+            
+            for det in detections:
+                x1, y1, x2, y2 = det.bbox
+                center_x = (x1 + x2) // 2
+                center_y = (y1 + y2) // 2
+                
+                safety_det = SafetyDetection(
+                    class_id=det.class_id,
+                    class_name=det.class_name,
+                    confidence=det.confidence,
+                    bbox=det.bbox,
+                    center=(center_x, center_y)
+                )
+                results.append(safety_det)
+                
+        except Exception as e:
+            print(f"RKNN 检测错误: {e}")
+        
+        return results
+    
+    def _detect_yolo(self, frame: np.ndarray) -> List[SafetyDetection]:
+        """使用 YOLO 模型检测"""
         results = []
         results = []
         
         
         try:
         try:
@@ -160,32 +483,26 @@ class SafetyDetector:
                     continue
                     continue
                 
                 
                 for i in range(len(boxes)):
                 for i in range(len(boxes)):
-                    # 获取类别
                     cls_id = int(boxes.cls[i])
                     cls_id = int(boxes.cls[i])
                     
                     
-                    # 只处理我们关心的类别
                     if cls_id not in self.CLASS_MAP:
                     if cls_id not in self.CLASS_MAP:
                         continue
                         continue
                     
                     
                     cls_name = self.CLASS_MAP[cls_id]
                     cls_name = self.CLASS_MAP[cls_id]
                     conf = float(boxes.conf[i])
                     conf = float(boxes.conf[i])
                     
                     
-                    # 根据类别设置不同的置信度阈值
                     threshold = self.person_threshold if cls_id == 3 else self.conf_threshold
                     threshold = self.person_threshold if cls_id == 3 else self.conf_threshold
                     if conf < threshold:
                     if conf < threshold:
                         continue
                         continue
                     
                     
-                    # 获取边界框
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     x1, y1, x2, y2 = map(int, xyxy)
                     x1, y1, x2, y2 = map(int, xyxy)
                     
                     
-                    # 过滤过小的检测框
                     width = x2 - x1
                     width = x2 - x1
                     height = y2 - y1
                     height = y2 - y1
                     if width < 10 or height < 10:
                     if width < 10 or height < 10:
                         continue
                         continue
                     
                     
-                    # 计算中心点
                     center_x = (x1 + x2) // 2
                     center_x = (x1 + x2) // 2
                     center_y = (y1 + y2) // 2
                     center_y = (y1 + y2) // 2
                     
                     
@@ -199,10 +516,17 @@ class SafetyDetector:
                     results.append(detection)
                     results.append(detection)
                     
                     
         except Exception as e:
         except Exception as e:
-            print(f"检测错误: {e}")
+            print(f"YOLO 检测错误: {e}")
         
         
         return results
         return results
     
     
+    def release(self):
+        """释放模型资源"""
+        if self.rknn_detector:
+            self.rknn_detector.release()
+            self.rknn_detector = None
+        self.model = None
+    
     def check_safety(self, frame: np.ndarray, 
     def check_safety(self, frame: np.ndarray, 
                      detections: List[SafetyDetection] = None) -> List[PersonSafetyStatus]:
                      detections: List[SafetyDetection] = None) -> List[PersonSafetyStatus]:
         """
         """
@@ -420,20 +744,23 @@ class LLMSafetyDetector:
     def __init__(self, yolo_model_path: str = None, 
     def __init__(self, yolo_model_path: str = None, 
                  llm_config: Dict[str, Any] = None,
                  llm_config: Dict[str, Any] = None,
                  use_gpu: bool = True,
                  use_gpu: bool = True,
-                 use_llm: bool = True):
+                 use_llm: bool = True,
+                 model_type: str = 'auto'):
         """
         """
         初始化检测器
         初始化检测器
         
         
         Args:
         Args:
-            yolo_model_path: YOLO 模型路径
+            yolo_model_path: 模型路径 (.pt, .rknn, .onnx)
             llm_config: 大模型配置
             llm_config: 大模型配置
-            use_gpu: 是否使用 GPU
+            use_gpu: 是否使用 GPU (仅 YOLO 模型有效)
             use_llm: 是否使用大模型判断
             use_llm: 是否使用大模型判断
+            model_type: 模型类型 ('auto', 'yolo', 'rknn', 'onnx')
         """
         """
-        # YOLO 检测器
+        # 安全检测器 (支持 YOLO/RKNN/ONNX)
         self.yolo_detector = SafetyDetector(
         self.yolo_detector = SafetyDetector(
             model_path=yolo_model_path,
             model_path=yolo_model_path,
-            use_gpu=use_gpu
+            use_gpu=use_gpu,
+            model_type=model_type
         )
         )
         
         
         # 大模型分析器
         # 大模型分析器

+ 3 - 2
dual_camera_system/safety_main.py

@@ -16,7 +16,7 @@ import argparse
 import logging
 import logging
 import threading
 import threading
 import signal
 import signal
-from typing import Optional
+from typing import Optional, List
 
 
 import cv2
 import cv2
 import numpy as np
 import numpy as np
@@ -180,7 +180,8 @@ class SafetyMonitorSystem:
                     yolo_model_path=self.config.get('model_path', SAFETY_DETECTION_CONFIG.get('model_path')),
                     yolo_model_path=self.config.get('model_path', SAFETY_DETECTION_CONFIG.get('model_path')),
                     llm_config=llm_config,
                     llm_config=llm_config,
                     use_gpu=self.config.get('use_gpu', SAFETY_DETECTION_CONFIG.get('use_gpu', True)),
                     use_gpu=self.config.get('use_gpu', SAFETY_DETECTION_CONFIG.get('use_gpu', True)),
-                    use_llm=self.enable_llm
+                    use_llm=self.enable_llm,
+                    model_type=self.config.get('model_type', SAFETY_DETECTION_CONFIG.get('model_type', 'auto'))
                 )
                 )
                 logger.info("安全检测器初始化成功")
                 logger.info("安全检测器初始化成功")
             except Exception as e:
             except Exception as e: