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 4 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 = {
-    'ip': '192.168.1.100',      # 全景摄像头IP
-    'port': 37777,              # 大华SDK默认端口
+    'ip': '192.168.20.196',      # 全景摄像头IP
+    'port': 554,                  # RTSP端口
     '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 = {
-    'ip': '192.168.1.101',      # 球机IP
-    'port': 37777,
+    'ip': '192.168.20.197',      # 球机IP
+    'port': 554,                 # RTSP端口
     '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_PATH = {
-    'lib_path': '/home/wen/dsh/dh/Bin',
+    'lib_path': '/Users/wenhongquan/Desktop/阿里云同步/项目/dnn/德胜河 AI/dsh/dh/Bin',
     'netsdk': 'libdhnetsdk.so',
 }

+ 8 - 2
dual_camera_system/config/detection.py

@@ -11,8 +11,14 @@ 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,           # 一般物品置信度阈值 (安全帽、反光衣)
     'person_threshold': 0.8,         # 人员检测置信度阈值
     

+ 4 - 4
dual_camera_system/config/system.py

@@ -24,12 +24,12 @@ SYSTEM_CONFIG = {
     'enable_ptz_tracking': True,         # 启用 PTZ 跟踪联动
     
     # 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'(编号识别)

+ 166 - 19
dual_camera_system/panorama_camera.py

@@ -8,6 +8,7 @@ import numpy as np
 import threading
 import queue
 import time
+import os
 from typing import Optional, List, Tuple, Dict, Any
 from dataclasses import dataclass
 
@@ -330,41 +331,95 @@ class ObjectDetector:
     """
     物体检测器
     使用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:
-            model_path: 模型路径 (自定义模型)
+            model_path: 模型路径 (支持 .pt, .rknn, .onnx)
             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.rknn_detector = None
         self.model_path = model_path
         self.use_gpu = use_gpu
         self.model_size = model_size
+        self.model_type = model_type
         self.config = DETECTION_CONFIG
         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()
     
     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检测模型"""
         try:
-            # 使用ultralytics YOLO11
             from ultralytics import YOLO
             
             if self.model_path:
-                # 使用自定义模型
                 self.model = YOLO(self.model_path)
             else:
-                # 使用YOLO11预训练模型
-                # YOLO11模型命名: yolo11n.pt, yolo11s.pt, yolo11m.pt, yolo11l.pt, yolo11x.pt
                 model_name = f'yolo11{self.model_size}.pt'
                 self.model = YOLO(model_name)
             
-            # 预热模型
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             self.model(dummy, device=self.device, verbose=False)
             
@@ -378,9 +433,95 @@ class ObjectDetector:
     
     def _load_opencv_model(self):
         """使用OpenCV加载模型"""
-        # 可以加载ONNX模型
         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]:
         """
         使用YOLO11检测物体
@@ -389,13 +530,19 @@ class ObjectDetector:
         Returns:
             检测结果列表
         """
-        if self.model is None or frame is None:
+        if frame is None:
             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 = []
         
         try:
-            # YOLO11推理
             detections = self.model(
                 frame, 
                 device=self.device, 
@@ -409,28 +556,22 @@ class ObjectDetector:
                     continue
                     
                 for i in range(len(boxes)):
-                    # 获取类别
                     cls_id = int(boxes.cls[i])
                     cls_name = det.names[cls_id]
                     
-                    # 过滤目标类别
                     if cls_name not in self.config['target_classes']:
                         continue
                     
-                    # 获取置信度
                     conf = float(boxes.conf[i])
                     
-                    # 获取边界框
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     x1, y1, x2, y2 = map(int, xyxy)
                     width = x2 - x1
                     height = y2 - y1
                     
-                    # 过滤过小的检测框
                     if width < 10 or height < 10:
                         continue
                     
-                    # 计算中心点
                     center_x = x1 + width // 2
                     center_y = y1 + height // 2
                     
@@ -455,8 +596,6 @@ class ObjectDetector:
         Returns:
             带关键点的检测结果列表
         """
-        # 如果使用pose模型,可以获取人体关键点
-        # 用于更精确的人体定位
         return self.detect(frame)
     
     def detect_persons(self, frame: np.ndarray) -> List[DetectedObject]:
@@ -469,6 +608,14 @@ class ObjectDetector:
         """
         results = self.detect(frame)
         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:

+ 352 - 25
dual_camera_system/safety_detector.py

@@ -2,6 +2,10 @@
 施工现场安全行为检测模块
 使用 YOLO11 模型检测人员、安全帽、反光衣
 判断是否存在违规行为(未戴安全帽、未穿反光衣)
+
+支持两种模型格式:
+- YOLO (.pt/.onnx): 使用 ultralytics 库
+- RKNN (.rknn): 使用 rknnlite 库 (RK3588 平台)
 """
 
 import cv2
@@ -9,6 +13,239 @@ import numpy as np
 from typing import Optional, List, Tuple, Dict, Any
 from dataclasses import dataclass
 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):
@@ -81,15 +318,12 @@ class SafetyDetector:
     使用 YOLO11 检测人员、安全帽、反光衣
     """
     
-    # 类别映射 (根据 yolo11m_safety.pt 模型的训练标签)
-    # 0: 安全帽, 3: 人, 4: 安全衣/反光衣
     CLASS_MAP = {
         0: '安全帽',
         3: '人',
         4: '反光衣'
     }
     
-    # 反向映射
     CLASS_ID_MAP = {
         'helmet': 0,
         'person': 3,
@@ -97,44 +331,94 @@ class SafetyDetector:
     }
     
     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:
-            model_path: 模型路径,默认使用 yolo11m_safety.pt
-            use_gpu: 是否使用 GPU
+            model_path: 模型路径,默认使用 yolo11m_safety.pt 或 .rknn
+            use_gpu: 是否使用 GPU (仅 YOLO 模型有效)
             conf_threshold: 一般物品置信度阈值 (安全帽、反光衣)
             person_threshold: 人员检测置信度阈值
+            model_type: 模型类型 ('auto', 'yolo', 'rknn', 'onnx')
         """
         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.device = 'cuda:0' if use_gpu else 'cpu'
         
-        # 置信度阈值
         self.conf_threshold = conf_threshold
         self.person_threshold = person_threshold
         
-        # 加载模型
         self._load_model()
     
     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 安全检测模型"""
         try:
             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)
             
-            # 预热模型
             dummy = np.zeros((640, 640, 3), dtype=np.uint8)
             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:
             raise ImportError("未安装 ultralytics,请运行: pip install ultralytics")
         except Exception as e:
-            raise RuntimeError(f"加载模型失败: {e}")
+            raise RuntimeError(f"加载 YOLO 模型失败: {e}")
     
     def detect(self, frame: np.ndarray) -> List[SafetyDetection]:
         """
@@ -146,9 +430,48 @@ class SafetyDetector:
         Returns:
             检测结果列表
         """
-        if self.model is None or frame is None:
+        if frame is None:
             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 = []
         
         try:
@@ -160,32 +483,26 @@ class SafetyDetector:
                     continue
                 
                 for i in range(len(boxes)):
-                    # 获取类别
                     cls_id = int(boxes.cls[i])
                     
-                    # 只处理我们关心的类别
                     if cls_id not in self.CLASS_MAP:
                         continue
                     
                     cls_name = self.CLASS_MAP[cls_id]
                     conf = float(boxes.conf[i])
                     
-                    # 根据类别设置不同的置信度阈值
                     threshold = self.person_threshold if cls_id == 3 else self.conf_threshold
                     if conf < threshold:
                         continue
                     
-                    # 获取边界框
                     xyxy = boxes.xyxy[i].cpu().numpy()
                     x1, y1, x2, y2 = map(int, xyxy)
                     
-                    # 过滤过小的检测框
                     width = x2 - x1
                     height = y2 - y1
                     if width < 10 or height < 10:
                         continue
                     
-                    # 计算中心点
                     center_x = (x1 + x2) // 2
                     center_y = (y1 + y2) // 2
                     
@@ -199,10 +516,17 @@ class SafetyDetector:
                     results.append(detection)
                     
         except Exception as e:
-            print(f"检测错误: {e}")
+            print(f"YOLO 检测错误: {e}")
         
         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, 
                      detections: List[SafetyDetection] = None) -> List[PersonSafetyStatus]:
         """
@@ -420,20 +744,23 @@ class LLMSafetyDetector:
     def __init__(self, yolo_model_path: str = None, 
                  llm_config: Dict[str, Any] = None,
                  use_gpu: bool = True,
-                 use_llm: bool = True):
+                 use_llm: bool = True,
+                 model_type: str = 'auto'):
         """
         初始化检测器
         
         Args:
-            yolo_model_path: YOLO 模型路径
+            yolo_model_path: 模型路径 (.pt, .rknn, .onnx)
             llm_config: 大模型配置
-            use_gpu: 是否使用 GPU
+            use_gpu: 是否使用 GPU (仅 YOLO 模型有效)
             use_llm: 是否使用大模型判断
+            model_type: 模型类型 ('auto', 'yolo', 'rknn', 'onnx')
         """
-        # YOLO 检测器
+        # 安全检测器 (支持 YOLO/RKNN/ONNX)
         self.yolo_detector = SafetyDetector(
             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 threading
 import signal
-from typing import Optional
+from typing import Optional, List
 
 import cv2
 import numpy as np
@@ -180,7 +180,8 @@ class SafetyMonitorSystem:
                     yolo_model_path=self.config.get('model_path', SAFETY_DETECTION_CONFIG.get('model_path')),
                     llm_config=llm_config,
                     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("安全检测器初始化成功")
             except Exception as e: