# 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