基于 YOLOv8 的反无人机红外检测系统构建与部署
使用 YOLOv8 训练一个基于无人机目标检测的数据集，包括反无人机红外数据集（Anti-UAV410 IR），从环境配置、数据准备、模型训练、推理实现等方面进行详细说明。以下是完整的步骤指南。
一、环境配置
1. 安装 CUDA 驱动
确保计算机上安装了与你的 NVIDIA GPU 兼容的 CUDA 版本。你可以从 NVIDIA 官网 下载并安装适合你操作系统的 CUDA Toolkit。
2. 安装 Anaconda
访问 Anaconda 官网，根据你的操作系统下载并安装 Anaconda。
数据集描述：
- yolo 无人机目标检测数据集
- 反无人机红外数据集 Anti-UAV410 IR
- yolo 格式 yolo 可直接训练的格式
- 训练集 213995 张 测试集 94711 张
3. 创建 Python 虚拟环境
打开终端或命令提示符，创建一个新的 Anaconda 环境：
conda create -n drone_detection python=3.9
conda activate drone_detection
4. 安装必要的依赖项
在激活的环境中安装必要的库：
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
pip install ultralytics opencv-python matplotlib numpy
二、数据准备
假设你的数据集已经按照以下结构组织，并且每个图像都有一个对应的 YOLO 格式的标签文件（.txt）：
drone_dataset/
├── images/
│ ├── train/
│ ├── val/
│ └── test/
└── labels/
├── train/
├── val/
└── test/
创建一个描述数据集路径和类别的 .yaml 文件（例如 data.yaml）：
train: ./drone_dataset/images/train
val: ./drone_dataset/images/val
test: ./drone_dataset/images/test
nc: 1
names: ['drone']
三、模型训练与评估
1. 训练代码
首先，下载官方预训练模型（YOLOv8s）：
wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt
然后，开始训练：
from ultralytics import YOLO
# 加载预训练模型
model = YOLO('yolov8s.pt')
# 开始训练
results = model.train(data='path/to/data.yaml', epochs=100, imgsz=640)
或者直接通过命令行执行：
yolo train data=path/to/data.yaml model=yolov8s.pt epochs=100 imgsz=640 batch=16 device=0
2. 推理代码
加载训练好的模型进行预测：
from ultralytics import YOLO
import cv2
# 加载训练好的模型
model = YOLO('runs/detect/train/weights/best.pt')
def detect_drone(image_path, model):
# 使用模型进行预测
results = model.predict(source=image_path, conf=0.25)
# 可调整置信度阈值
# 获取原始图像并绘制检测框
for result in results:
annotated_img = result.plot()
# 自动绘制边界框和标签
annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
# 转换颜色空间
return annotated_img
# 示例调用并显示图像
result_image = detect_drone('test_images/drone.jpg', model)
cv2.imshow('Drone Detection Result', result_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
3. 评估代码
在验证集上评估模型性能：
from ultralytics import YOLO
# 加载训练好的模型
model = YOLO('runs/detect/train/weights/best.pt')
# 在验证集上评估模型性能
metrics = model.val(data='path/to/data.yaml')
print("mAP50:", metrics.box.map50)
print("mAP50-95:", metrics.box.map)
print("Precision:", metrics.box.precision)
print("Recall:", metrics.box.recall)
四、核心功能实现
你已经训练好了 YOLOv8 的无人机检测模型权重（如 best.pt），现在希望根据这些权重构建一个完整的无人机目标检测系统，包括实时检测、批量检测、可视化及结果保存等。
1. 项目结构建议
drone_detection_system/
├── weights/
│ └── best.pt
├── detect.py
├── utils/
│ └── plotting.py
├── input/
│ └── test.mp4
├── output/
│ └── result.mp4
├── config/
│ └── data.yaml
└── requirements.txt
2. 环境依赖
pip install ultralytics opencv-python numpy matplotlib
3. 加载模型（detect.py）
from ultralytics import YOLO
# 加载本地训练好的模型
model = YOLO('weights/best.pt')
4. 图像检测函数
import cv2
def detect_image(image_path, model, save_path=None):
results = model(image_path)
for r in results:
im_array = r.plot()
im = cv2.cvtColor(im_array, cv2.COLOR_RGB2BGR)
if save_path:
cv2.imwrite(save_path, im)
cv2.imshow("Detection", im)
cv2.waitKey(0)
cv2.destroyAllWindows()
5. 视频检测函数（支持摄像头、视频文件、RTSP）
def detect_video(video_source, model, save_path=None):
cap = cv2.VideoCapture(video_source)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(save_path, fourcc, fps, (w, h)) if save_path else None
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame)
for r in results:
annotated_frame = r.plot()
annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
if out:
out.write(annotated_frame)
cv2.imshow('Drone Detection', annotated_frame)
if cv2.waitKey(1) == 27:
break
cap.release()
if out:
out.release()
cv2.destroyAllWindows()
6. 使用示例
if __name__ == '__main__':
model = YOLO('weights/best.pt')
# 图像检测
detect_image('input/test.jpg', model, 'output/result.jpg')
# 视频检测
detect_video('input/test.mp4', model, 'output/result.mp4')
# 实时摄像头检测
detect_video(0, model, None)
# RTSP 流检测
detect_video('rtsp://admin:[email protected]:554', model, None)
五、部署方案建议
1. 本地桌面系统（Windows / Linux）
- 使用 PyQt / Tkinter 构建 GUI 界面
- 支持图像、视频、摄像头、RTSP 检测
- 支持结果保存与导出
2. Web 系统（Flask + YOLO）
- 使用 Flask 构建 Web 接口
- 前端上传图像/视频，后端返回检测结果
- 可用于远程无人机监控平台
3. 边缘设备部署（Jetson Nano / Xavier NX）
- 使用 TensorRT 加速推理
- 部署
.engine模型文件 - 支持低功耗实时检测
4. Docker 容器化部署
FROM python:3.9-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["python", "detect.py"]
构建并运行：
docker build -t drone-detection .
docker run -it --rm --gpus all drone-detection
六、总结
| 功能 | 实现方式 |
|---|---|
| 模型加载 | YOLO('weights/best.pt') |
| 图像检测 | model.predict() + OpenCV 绘图 |
| 视频检测 | cv2.VideoCapture + 循环推理 |
| 结果保存 | 图像、视频、JSON 格式 |
| 部署方式 | 本地、Web、Docker、嵌入式设备 |
附：完整 detect.py 示例
from ultralytics import YOLO
import cv2
import json
model = YOLO('weights/best.pt')
def detect_image(image_path, save_path=None):
results = model(image_path)
for r in results:
im_array = r.plot()
im = cv2.cvtColor(im_array, cv2.COLOR_RGB2BGR)
if save_path:
cv2.imwrite(save_path, im)
cv2.imshow("Detection", im)
cv2.waitKey(0)
cv2.destroyAllWindows()
def detect_video(video_source, save_path=None):
cap = cv2.VideoCapture(video_source)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(save_path, fourcc, fps, (w, h)) if save_path else None
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame)
for r in results:
annotated_frame = r.plot()
annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
if out:
out.write(annotated_frame)
cv2.imshow('Drone Detection', annotated_frame)
if cv2.waitKey(1) == 27:
break
cap.release()
if out:
out.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
detect_video(, )
