289 lines
7.3 KiB
Markdown
289 lines
7.3 KiB
Markdown
# 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行
|
||
|
||
**函数签名**:
|
||
```python
|
||
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（默认正常）
|
||
|
||
**推荐扩展方式**:
|
||
```python
|
||
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 | 原始帧队列大小 |
|
||
|
||
---
|
||
|
||
## 🔧 自定义算法开发
|
||
|
||
### 添加新的检测算法
|
||
|
||
```python
|
||
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 |