# UAV Detection System

基于 ONNX 的无人机检测系统，支持单张图片和批量处理。

## 概述

本项目实现了多种UAV（无人机）检测模型的推理系统，支持三种不同的模型架构和相应的后处理方法。每种模型都有其特定的输入输出格式、后处理流程和可视化方式。

## 功能特点

- 支持 ONNX 模型推理
- 支持 CUDA 加速
- 支持批量处理图片
- 自动生成检测报告（CSV格式）
- 支持检测框面积比例限制
- 支持保存未检测到目标的图片
- 支持三种不同架构的模型自动识别
- 集成边界框验证机制

## 支持的模型类型

### 1. Anti_UAV 模型
- **文件名**: `250411_Anti_UAV.onnx`
- **识别特征**: 输入包含 `scale_factor`，输出名称以 `multiclass_nms3` 开头
- **输入尺寸**: 640×640
- **类别**: 单类别检测（UAV）

### 2. UAV-250411 模型
- **文件名**: `UAV-250411.onnx`
- **识别特征**: 输入包含 `scale_factor`，输出为2个张量且第一个以 `tmp_` 开头
- **输入尺寸**: 640×640
- **类别**: 单类别检测（UAV）

### 3. uav_and_bird 模型
- **文件名**: `uav_and_bird.onnx`
- **识别特征**: 输入不包含 `scale_factor`，输出张量数量大于2
- **输入尺寸**: 640×640
- **类别**: 双类别检测（Bird: 0, UAV: 1）

## 环境要求

- Python 3.8+
- OpenCV 4.5+
- ONNX Runtime 1.9+
- CUDA 11.0+ (可选，用于GPU加速)

## 安装依赖

```bash
pip install -r requirements.txt
```

## 使用方法

### 命令行模式

```bash
# 处理单张图片
python -m src.core.inference --input path/to/image.jpg --output results

# 处理整个目录
python -m src.core.inference --input path/to/images_dir --output results

# 调整检测参数
python -m src.core.inference --input path/to/image.jpg --threshold 0.6 --max-bbox-ratio 0.1

# 保存未检测到目标的图片
python -m src.core.inference --input path/to/images_dir --save-empty

# 自动识别模型类型
python inference.py --input /path/to/images --threshold 0.5

# 指定模型类型
python inference.py --input /path/to/images --model-type uav_and_bird

# 批量处理
python inference.py --input /path/to/directory --output /path/to/results
```

### 图形界面模式

```bash
python -m src.core.inference --gui
```

### 编程接口

```python
# 初始化检测器
detector = ONNXDetector(
    input_dir="/path/to/images",
    model_type="uav_and_bird",
    prob_threshold=0.5
)

# 处理单张图像
detections, img_out, detection_list = detector.process_image("/path/to/image.jpg")

# 批量处理
total_detections = detector.process_directory()
```

## 参数说明

- `--input`: 输入图像路径或目录（必需）
- `--output`: 输出目录路径，默认为输入目录名+_results
- `--threshold`: 检测置信度阈值，默认0.5
- `--max-bbox-ratio`: 检测框最大面积比例阈值，默认0.05
- `--save-empty`: 是否保存未检测到目标的图片
- `--gui`: 启用图形界面选择输入目录
- `--model-type`: 指定模型类型（可选，系统可自动识别）

## 模型结构分析

### 输入预处理

所有模型都采用相同的预处理流程：

```python
# 图像预处理步骤
1. 颜色空间转换: BGR → RGB
2. 尺寸调整: 原图 → 640×640
3. 数据类型转换: uint8 → float32
4. 归一化: [0,255] → [0,1]
5. 标准化: 使用ImageNet均值和标准差
   - mean = [0.485, 0.456, 0.406]
   - std = [0.229, 0.224, 0.225]
6. 维度调整: HWC → CHW
7. 批次维度: CHW → NCHW
```

### 模型输入格式

#### Anti_UAV 和 UAV-250411 模型
```python
inputs = {
    'image': img[None, :, :, :],        # (1, 3, 640, 640)
    'scale_factor': scale_factor[None, :] # (1, 2)
}
```

#### uav_and_bird 模型
```python
inputs = {
    'images': img[None, :, :, :]  # (1, 3, 640, 640)
}
```

## 后处理方法详解

### 1. Anti_UAV 模型后处理

**特点**: 模型内置NMS，输出已经过滤的检测结果

```python
# 输出格式
bbox: (N, 4)      # 边界框坐标 [x1, y1, x2, y2]
confidence: (N, 1) # 置信度分数

# 后处理流程
1. 置信度过滤: confidence > threshold
2. 坐标缩放: 模型输出 → 原图尺寸
3. 边界框验证: 面积比例检查
4. 结果保存和可视化
```

### 2. UAV-250411 模型后处理

**特点**: 需要手动实现NMS处理

```python
# 输出格式
output[0]: (N, 4)  # 边界框坐标
output[1]: (N, 1)  # 置信度分数

# 后处理流程
1. 置信度过滤: confidence > threshold
2. 坐标缩放: 模型输出 → 原图尺寸
3. NMS处理: IoU阈值 = 0.4
4. 边界框验证
5. 结果保存和可视化
```

### 3. uav_and_bird 模型后处理

**特点**: 多尺度输出，使用Distribution Focal Loss，需要复杂的解码过程

```python
# 输出格式（3个尺度）
small_output: (1, 80, 80, 21)   # 小目标检测
medium_output: (1, 40, 40, 21)  # 中等目标检测
large_output: (1, 20, 20, 21)   # 大目标检测

# 每个输出的通道分布
# 21 = 17(bbox) + 2(classes) + 2(quality)

# 后处理流程
1. 多尺度处理:
   - 小尺度: stride=8, 网格80×80
   - 中尺度: stride=16, 网格40×40  
   - 大尺度: stride=32, 网格20×20

2. Distribution Focal Loss解码:
   - bbox_pred: 17维距离分布 → 4个距离值
   - 使用softmax + 期望值计算实际距离

3. 网格坐标生成（带0.5偏移）:
   grid_x, grid_y = np.meshgrid(range(W), range(H))
   grid_x = (grid_x.flatten() + 0.5)
   grid_y = (grid_y.flatten() + 0.5)

4. 边界框计算:
   x1 = (grid_x - d_left) * stride
   y1 = (grid_y - d_top) * stride
   x2 = (grid_x + d_right) * stride
   y2 = (grid_y + d_bottom) * stride

5. 置信度计算:
   final_scores = cls_pred * quality_pred

6. 多尺度结果合并和NMS处理
```

## 可视化和画框方式

### 1. 边界框绘制

所有模型都使用相同的绘制方式：

```python
# 边界框样式
cv2.rectangle(img_out, (x1, y1), (x2, y2), (255, 0, 0), 4)  # 蓝色框，线宽4

# 标签样式
label = f'{class_name} {confidence:.2f}'
cv2.putText(img_out, label, (x1, y1 - 10), 
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)  # 红色文字
```

### 2. 类别标签映射

```python
# Anti_UAV 和 UAV-250411 模型
class_names = {0: 'UAV'}

# uav_and_bird 模型
class_names = {0: 'Bird', 1: 'UAV'}
```

### 3. 检测结果保存

```python
# 保存检测到的目标区域
target_filename = f"{base_name}_{detection_count}.jpg"
cv2.imwrite(os.path.join(targets_dir, target_filename), roi)

# 保存带标注的完整图像
output_filename = f"detected_{base_name}"
cv2.imwrite(os.path.join(output_dir, output_filename), img_out)
```

## 性能优化和质量控制

### 1. 边界框验证

```python
# 面积比例检查
bbox_area = (x2 - x1) * (y2 - y1)
image_area = orig_w * orig_h
if bbox_area / image_area > max_bbox_ratio:  # 默认0.05
    continue

# 尺寸有效性检查
if x2 <= x1 or y2 <= y1:
    continue
```

### 2. NMS参数

```python
# 所有模型统一使用的NMS参数
conf_threshold = 0.5    # 置信度阈值
iou_threshold = 0.4     # IoU阈值
```

## 模型自动识别机制

系统通过分析ONNX模型的输入输出结构自动识别模型类型：

```python
def get_model_type(model_path: str) -> str:
    model = onnx.load(model_path)
    input_names = [i.name for i in model.graph.input]
    output_names = [o.name for o in model.graph.output]
    
    if 'scale_factor' in input_names:
        if any(name.startswith('multiclass_nms3') for name in output_names):
            return 'Anti_UAV'
        else:
            return 'UAV-250411'
    else:
        if len(output_names) > 2:
            return 'uav_and_bird'
    
    return 'unknown'
```

## 输出说明

程序会在输出目录中生成以下内容：

- 检测结果图片（带检测框）
- 检测到的目标区域图片
- `detection_report.csv`: 检测报告，包含以下信息：
  - 图片路径
  - 检测时间
  - 是否检测到目标
  - 检测框坐标
  - 置信度分数

## 注意事项

1. 确保模型文件路径正确
2. 如果使用GPU加速，请确保CUDA环境配置正确
3. 批量处理时建议使用相对较小的图片尺寸以提高处理速度
4. 检测报告会自动覆盖同名文件，请注意备份
5. 系统会自动识别模型类型，无需手动指定
6. 支持多种模型架构，可根据实际需求选择合适的模型

## 项目特点总结

本项目实现了一个灵活的多模型UAV检测系统，具有以下特点：

1. **模型兼容性**: 支持三种不同架构的ONNX模型
2. **自动识别**: 根据模型结构自动选择对应的后处理方法
3. **质量控制**: 集成边界框验证机制
4. **可扩展性**: 易于添加新的模型类型和后处理方法
5. **性能优化**: 支持CUDA加速和批量处理
6. **用户友好**: 支持命令行和图形界面两种使用方式
7. **完整输出**: 自动生成检测报告和可视化结果

每种模型都有其特定的应用场景和优势，用户可以根据实际需求选择合适的模型进行部署。

## 更新日志

### 2024-03-29
- 移除Excel报告生成功能
- 优化检测报告生成逻辑
- 修复模型加载和属性访问问题
- 移除ORB特征匹配误报过滤功能
- 清理调试信息，优化代码结构
- 合并技术文档到README，提供完整的使用和技术说明