旋转框 YOLO 训练代码（YOLOv8-OBB）如何训练无人机视角下多模态红外可见光红外对齐车辆目标检测数据集 多模态目标检测 （RGB + IR 融合） - 旋转目标检测 （带方向角的车辆定位）
无人机视角下RGB+红外对齐车辆目标检测数据集
模态与视角：无人机搭载双光相机，同时采集 RGB + 红外，成对图像，对应同一视角的交通场景。 规模与分辨率：共 28,439 对 RGB-IR 图像（56,878 张），图像尺寸约 840×712 标注与类别：五类车辆目标——car、truck、bus、van、freight car，提供带方向角的 bbox，适合做多模态检测与旋转框检测
。1
1
“无人机视角下RGB+红外对齐车辆目标检测数据集”的结构化表格描述：
| 属性类别 | 详细描述 |
|---|---|
| 数据集名称 | UAV-RGB-IR-Vehicle（无人机双光车辆检测数据集） |
| 采集平台 | 无人机搭载同步双光相机（可见光 + 热红外） |
| 成像模态 | 成对图像： • RGB（可见光） • IR（热红外） → 严格时空对齐，同一视角、同一时刻 |
| 场景类型 | 交通道路场景（城市道路、高速、交叉口等），低空航拍（典型 UAV 视角） |
| 图像总数 | 28,439 对（共 56,878 张 单模态图像） |
| 图像分辨率 | 约 840 × 712 像素（每对 RGB 与 IR 分辨率一致） |
| 目标类别 | 共 5 类车辆： 1. car（小轿车）2. truck（卡车）3. bus（公交车）4. van（厢式货车）5. freight car（货运车/重型载具） |
| 标注形式 | - 旋转边界框（Rotated Bounding Box） 格式： (x_center, y_center, width, height, angle)- 每个目标包含 方向角（orientation angle），支持朝向感知检测 - 所有标注在 RGB 与 IR 图像上共享同一坐标系（已对齐） |
| 标注数量 | 未明确给出总实例数，但覆盖密集交通场景，含遮挡、小目标、多尺度目标 |
| 文件组织建议 | <br>dataset/<br>├── rgb/<br>│ ├── 000001.jpg<br>│ └── ...<br>├── ir/<br>│ ├── 000001.jpg<br>│ └── ...<br>└── labels/<br> ├── 000001.txt<br> └── ...<br>每个 .txt 文件对应一对图像的旋转框标注 |
| 标注格式示例 | 每行：<class_id> <x_center> <y_center> <width> <height> <angle>（坐标归一化至 [0,1]，角度单位：弧度 或 度，需注明） |
| 适用任务 | - 多模态目标检测（RGB + IR 融合） - 旋转目标检测（带方向角的车辆定位） - 跨模态对齐学习 - 夜间/低照度鲁棒检测（利用 IR 模态） - 小目标与密集车辆检测 |
| 数据特点 | - 严格对齐的双模态数据 - 真实交通场景复杂性（光照变化、遮挡、尺度差异） - 提供方向信息，适用于自动驾驶、交通流分析等下游任务 - 大规模（>28k 对），适合深度学习训练 |
✅ 该数据集填补了无人机双光（RGB+IR）车辆检测领域缺乏带旋转框标注的大规模公开数据集的空白，特别适合研究：多模态特征融合机制旋转检测模型（如 R3Det、KLD-based Rotated RetinaNet）昼夜一致性检测系统
针对旋转框目标检测任务，YOLOv8 原版并不直接支持旋转边界框（Rotated Bounding Box, OBB）。为了实现这一功能，我们可以基于 YOLOv8 进行一些修改或使用已有的扩展版本，如 RoI-Transformer 或者 R3Det 等。这里提供一个简化的方案，展示如何基于 YOLOv8 的框架进行自定义以支持旋转框的训练。
一、前提条件
确保你已经安装了必要的依赖库：
pip install ultralytics opencv-python numpy tqdm 二、数据准备
假设你的数据集按以下结构组织，并且标注格式为 x_center, y_center, width, height, angle （角度单位可以是弧度或度数）：
UAV-RGB-IR-Vehicle/ ├── rgb/ │ ├── train/ │ ├── val/ │ └── test/ ├── ir/ │ ├── train/ │ ├── val/ │ └── test/ └── labels/ ├── train/ ├── val/ └── test/ 每个 .txt 文件包含每张图片的旋转框标注信息，格式如下：
<class_id> <x_center> <y_center> <width> <height> <angle> 三、配置文件 (uav_rgb_ir_vehicle.yaml)
创建一个配置文件来指定数据路径、类别等信息：
# uav_rgb_ir_vehicle.yamltrain: ./UAV-RGB-IR-Vehicle/rgb/train val: ./UAV-RGB-IR-Vehicle/rgb/val nc:5# 类别数量：car, truck, bus, van, freight carnames:- car - truck - bus - van - freight car # 根据需要调整其他参数四、修改 YOLOv8 模型以支持旋转框
由于 YOLOv8 不直接支持旋转框，我们需要对模型和损失函数做一些修改。这里给出一个简化版的示例，主要集中在损失计算部分。你可以基于现有的 YOLOv8 模型进行修改，或者寻找已有支持旋转框的模型架构作为基础。
修改后的损失函数（obb_loss.py）
import torch import torch.nn as nn classRotatedIoULoss(nn.Module):def__init__(self):super(RotatedIoULoss, self).__init__()defforward(self, pred_boxes, target_boxes):""" pred_boxes: [N, 5] tensor (x_center, y_center, width, height, angle) target_boxes: [N, 5] tensor (x_center, y_center, width, height, angle) """ pred_boxes = pred_boxes.clone().detach() target_boxes = target_boxes.clone().detach() pred_boxes[:,4]= pred_boxes[:,4]%180# 角度归一化到[0, 180) target_boxes[:,4]= target_boxes[:,4]%180# 计算旋转矩形的四个顶点坐标defget_vertices(boxes): vertices =[]for box in boxes: x_center, y_center, w, h, angle = box.tolist() angle_rad = angle *3.1415926/180.0 cos_a, sin_a = torch.cos(angle_rad), torch.sin(angle_rad) rect = torch.tensor([[-w /2,-h /2],[w /2,-h /2],[w /2, h /2],[-w /2, h /2]]) rotation_matrix = torch.tensor([[cos_a,-sin_a],[sin_a, cos_a]]) rotated_rect = torch.mm(rect, rotation_matrix) rotated_rect += torch.tensor([x_center, y_center]) vertices.append(rotated_rect)return torch.stack(vertices) pred_vertices = get_vertices(pred_boxes) target_vertices = get_vertices(target_boxes)# 计算交并比（IoU）defcompute_iou(polygon1, polygon2):from shapely.geometry import Polygon poly1 = Polygon(polygon1.cpu().numpy()) poly2 = Polygon(polygon2.cpu().numpy()) iou = poly1.intersection(poly2).area / poly1.union(poly2).area return iou ious = torch.tensor([compute_iou(p, t)for p, t inzip(pred_vertices, target_vertices)]) loss =1- ious.mean()return loss 五、训练脚本 (train_obb.py)
from ultralytics import YOLO import os from obb_loss import RotatedIoULoss defmain():# 创建输出目录 os.makedirs("runs/uav_rgb_ir_vehicle", exist_ok=True)# 加载预训练模型（建议使用 yolov8s 或 yolov8m） model = YOLO('yolov8s.pt')# 可替换为 'yolov8m.pt' / 'yolov8l.pt'# 自定义损失函数 custom_loss = RotatedIoULoss()# 开始训练 results = model.train( data='uav_rgb_ir_vehicle.yaml',# 数据配置文件 epochs=100,# 根据收敛情况调整 imgsz=840,# 图像尺寸（与数据一致） batch=16,# 根据 GPU 显存调整 name='yolov8s_uav_rgb_ir_vehicle', project='runs/uav_rgb_ir_vehicle', device=0,# GPU ID，多卡可用 [0,1] workers=8, cache=False,# 若内存充足可设为 True 加速 optimizer='AdamW', lr0=0.01,# 初始学习率 lrf=0.01,# 最终学习率 = lr0 * lrf momentum=0.937, weight_decay=0.0005, warmup_epochs=3, patience=30,# 早停：验证损失不再下降则停止 save=True, save_period=10,# 每10个epoch保存一次 verbose=True, plots=True,# 生成训练曲线图 loss_fn=custom_loss # 使用自定义的旋转框损失函数)print(f"✅ 训练完成！最佳模型路径：{results.save_dir}/weights/best.pt")if __name__ =='__main__': main()
