R360-UAVmodelTool

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加速)

安装依赖

pip install -r requirements.txt

使用方法

命令行模式

# 处理单张图片
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

图形界面模式

python -m src.core.inference --gui

编程接口

# 初始化检测器
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: 指定模型类型（可选，系统可自动识别）

模型结构分析

输入预处理

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

# 图像预处理步骤
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 模型

inputs = {
    'image': img[None, :, :, :],        # (1, 3, 640, 640)
    'scale_factor': scale_factor[None, :] # (1, 2)
}

uav_and_bird 模型

inputs = {
    'images': img[None, :, :, :]  # (1, 3, 640, 640)
}

后处理方法详解

1. Anti_UAV 模型后处理

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

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

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

2. UAV-250411 模型后处理

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

# 输出格式
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，需要复杂的解码过程

# 输出格式（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. 边界框绘制

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

# 边界框样式
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. 类别标签映射

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

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

3. 检测结果保存

# 保存检测到的目标区域
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. 边界框验证

# 面积比例检查
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参数

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

模型自动识别机制

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

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，提供完整的使用和技术说明