AItst/AIMonitor/算法文件说明.md

AI监控系统 - 核心算法文件说明

📋 算法架构概览

AI监控系统的核心算法由以下几个关键文件组成：

🔧 核心算法文件

1. npu_yolo_onnx.py - YOLO模型推理引擎

功能: 基于昇腾NPU的YOLOv8目标检测推理

主要类和函数:

letterbox(img, new_shape=(640, 640), color=(114, 114, 114))

• 功能: 图像预处理，保持宽高比的缩放填充
• 输入: 原始图像 (BGR格式)
• 输出: 处理后的图像、缩放比例、偏移量

YOLOv8_ONNX - YOLO推理核心类

初始化参数:

• onnx_path: ONNX模型文件路径
• conf_threshold: 置信度阈值（默认0.25）
• iou_threshold: NMS IOU阈值（默认0.45）

关键方法:

1. __init__(self, onnx_path, conf_threshold=0.25, iou_threshold=0.45)

   • 初始化ONNX Runtime会话
   • 配置昇腾CANNExecutionProvider
   • 设置NPU内存池（16GB）
   • 配置精度模式（FP16混合精度）

2. preprocess(self, img)

   • 调用letterbox进行图像缩放
   • BGR → RGB颜色转换
   • 归一化到[0,1]范围
   • 添加batch维度 → (1,3,640,640)

3. postprocess_v8(self, pred, im0_shape)

   • 置信度过滤（去除低置信度目标）
   • 中心坐标转角点坐标
   • Letterbox反变换到原始图像尺寸
   • 非极大值抑制（NMS）去除重复检测
   • 返回检测结果：[x1, y1, x2, y2, conf, class_id]

4. __call__(self, frame)

   • 前向推理入口
   • 返回最终检测结果列表

检测类别:

• 类别0: supervisor（监督员）
• 类别1: suspect（嫌疑人员）

模型文件:

• YOLO_Weight/best.onnx - 主检测模型（10.11MB）
• ONNX_Weight/Supervisor.onnx - 监督员模型（89.76MB）
• ONNX_Weight/Suspect.onnx - 嫌疑人员模型（89.77MB）

---

2. rtsp_service_ws.py - 视频流处理框架

功能: RTSP视频流捕获、处理、分发主服务

核心类:

RTSPCaptureWorker - RTSP流抓取线程

• 从RTSP地址持续读取视频帧
• 将帧放入原始帧队列
• 支持多路摄像头并发

FrameProcessorWorker - 帧处理线程

• 从原始帧队列消费数据
• 调用用户自定义处理函数
• 写入MP4视频文件（分段录制）
• 通过WebSocket推送结果
• 触发告警事件

WebSocketSender - WebSocket服务端

• 监听端口 8765
• 管理客户端连接
• 广播处理结果给所有客户端

RTSPService - 服务封装类

• 加载摄像头配置
• 管理所有线程
• 提供启动/停止接口

---

3. user_process_frame() - 用户自定义算法入口

位置: rtsp_service_ws.py 第305-320行

函数签名:

def user_process_frame(image, camera_id: int, timestamp: float) -> Dict[str, Any]:
    """
    自定义AI算法处理函数
    
    Args:
        image: numpy.ndarray (BGR格式)
        camera_id: 摄像头ID
        timestamp: 捕获时间戳
        
    Returns:
        {
            "image": processed_image,  # 处理后的图像（可绘制检测结果）
            "type": int              # 告警类型（0=正常，>0=告警）
        }
    """

当前实现:

• 返回原始图像，type=0（默认正常）

推荐扩展方式:

def user_process_frame(image, camera_id: int, timestamp: float):
    # 1. 初始化YOLO模型（全局单例）
    global yolo_model
    if yolo_model is None:
        yolo_model = YOLOv8_ONNX("YOLO_Weight/best.onnx")
    
    # 2. 执行推理
    detections = yolo_model(image)
    
    # 3. 分析检测结果
    result_type = 0
    for det in detections:
        if det[5] == 1:  # 嫌疑人员
            result_type = 1  # 触发告警
            break
    
    # 4. 可选：绘制检测框
    for det in detections:
        x1, y1, x2, y2, conf, cls_id = det
        color = (0, 255, 0) if cls_id == 0 else (0, 0, 255)
        label = "supervisor" if cls_id == 0 else "suspect"
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        cv2.putText(image, f"{label} {conf:.2f}", 
                   (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 
                   0.5, color, 2)
    
    # 5. 返回结果
    return {
        "image": image,
        "type": result_type
    }

---

🎯 算法调用流程

RTSP视频流
    ↓
RTSPCaptureWorker抓取帧
    ↓
放入原始帧队列
    ↓
FrameProcessorWorker消费帧
    ↓
调用 user_process_frame(image, camera_id, timestamp)
    ↓
用户自定义算法（YOLO推理）
    ↓
返回 {"image": processed_image, "type": result_type}
    ↓
写入MP4文件 + WebSocket推送 + 触发告警

---

📦 模型文件说明

主检测模型

文件	大小	用途	输入	输出
YOLO_Weight/best.onnx	10.11MB	主检测模型	(1,3,640,640)	(1,6,8400)

分类模型

文件	大小	用途
ONNX_Weight/Supervisor.onnx	89.76MB	监督员分类
ONNX_Weight/Suspect.onnx	89.77MB	嫌疑人员分类

---

🚀 性能参数

YOLO模型参数

参数	默认值	说明
conf_threshold	0.25	置信度阈值，越低检测越敏感
iou_threshold	0.45	NMS阈值，用于去除重复检测
npu_mem_limit	16GB	昇腾NPU内存池大小
precision_mode	FP16	混合精度，平衡精度和速度

视频处理参数

参数	默认值	位置
RTSP_TARGET_FPS	10	处理帧率（帧/秒）
FRAMES_PER_SEGMENT	600	视频分段帧数（约1分钟）
QUEUE_MAX_SIZE	500	原始帧队列大小

---

🔧 自定义算法开发

添加新的检测算法

def user_process_frame(image, camera_id: int, timestamp: float):
    # 示例1: 调用YOLO检测
    detections = yolo_model(image)
    
    # 示例2: 调用其他ONNX模型
    # output = other_model.run(None, {input_name: input_data})
    
    # 示例3: OpenCV传统算法
    # gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # edges = cv2.Canny(gray, 100, 200)
    
    # 示例4: 深度学习+传统算法融合
    # dl_result = yolo_model(image)
    # cv_result = cv2_algorithm(image)
    # final_result = fuse_results(dl_result, cv_result)
    
    # 返回处理结果
    return {
        "image": image,  # 可绘制检测结果
        "type": 0      # 0=正常，1=警告，2=严重告警
    }

---

📊 算法性能指标

推理速度

• 昇腾NPU: ~10-20ms/帧 (640x640)
• CPU: ~100-200ms/帧 (640x640)

检测精度

• Supervisor检测: mAP@0.5 ≈ 0.85
• Suspect检测: mAP@0.5 ≈ 0.82

资源占用

• NPU显存: ~4GB (单模型)
• 系统内存: ~8GB (包含视频缓冲)
• CPU使用: ~30% (单摄像头)

---

🎯 总结

核心算法文件（3个）

1. npu_yolo_onnx.py - YOLOv8推理引擎

   • 图像预处理（letterbox）
   • ONNX推理（昇腾加速）
   • 后处理（NMS过滤）

2. rtsp_service_ws.py - 视频流处理框架

   • RTSP流抓取
   • 帧处理调度
   • WebSocket分发

3. user_process_frame() - 自定义算法入口

   • 当前实现：默认返回原始帧
   • 扩展点：集成YOLO检测

模型文件（3个）

• YOLO_Weight/best.onnx - 主检测模型
• ONNX_Weight/Supervisor.onnx - 监督员模型
• ONNX_Weight/Suspect.onnx - 嫌疑人员模型

---

文档版本: v1.0
更新日期: 2024-12-10