MMRotate实战:从零构建自定义旋转目标检测数据集与模型训练

MMRotate实战:从零构建自定义旋转目标检测数据集与模型训练

1. 项目缘起:当通用检测框遇上旋转目标

在计算机视觉的日常开发中,我们经常会遇到一些“特殊”的目标。比如,遥感图像中倾斜停放的车辆、文档扫描件中带有角度的文字、或者工业场景中任意方向排列的零件。如果你尝试用经典的YOLO、Faster R-CNN等水平框检测器去处理这些目标,结果往往会让你哭笑不得——一个巨大的水平矩形框,把目标连同大量无关的背景一起框了进去,不仅定位不准,后续的IoU计算、NMS(非极大值抑制)都会出现问题,严重影响检测精度。

这就是旋转目标检测(Rotated Object Detection)要解决的问题。而MMRotate,正是基于强大的PyTorch框架和MMDetection、MMCV生态,为旋转目标检测量身打造的一个开源工具箱。它不是一个全新的算法,而是一个集成了多种主流旋转检测算法(如Rotated Faster R-CNN, Rotated RetinaNet, RoI Transformer等)的“脚手架”和“实验场”,让你能快速在自己的数据上验证和部署旋转检测模型。

我最近的一个项目恰好涉及对航拍图像中的建筑物进行精细定位,传统的水平框完全无法满足需求,于是便深入折腾了一番MMRotate。整个过程从数据准备、格式转换、模型训练到调优部署,踩了不少坑,也积累了一些实战心得。这篇文章,我就以一个“自定义数据集”的完整流程为线索,带你走通MMRotate,并分享那些官方文档里可能不会细说的细节。

2. 理解核心:旋转框的表示法与数据集构建

在动手写代码之前,我们必须先统一“语言”。旋转框怎么描述?你的数据标签又该如何准备?这是所有后续工作的基石,理解偏差会导致训练完全失败。

2.1 旋转框的几种表示法

旋转矩形框通常由一个中心点、宽高和旋转角度来定义。但就是这个“角度”,在不同的库和论文中有着不同的约定,MMRotate主要支持两种:

  1. OpenCV 表示法 (默认):(x_center, y_center, width, height, angle)

    • angle的单位是度数。
    • 角度范围(-90, 0]。这是最容易混淆的一点。在这种表示下,width被定义为旋转框边界框(Bounding Box)在旋转后,其自身坐标系下的长边。而angle表示的是这个width边相对于水平轴(x轴)顺时针旋转的角度,且被限制在-90度到0度之间(即-90 < angle <= 0)。可以理解为,框总是“站立”的,其较长的边(width)与水平线的夹角在-90度到0度这个区间内。这种表示法能唯一确定一个矩形(不考虑180度翻转的歧义)。
  2. 长边定义法: 另一种常见表示是(x_center, y_center, width, height, angle),但这里的width永远是旋转框自身的长边height是短边,angle是长边与x轴的夹角,范围通常是[-90, 90)[0, 180)。MMRotate的部分算法也支持这种模式,但需要在配置文件中明确指定angle_version

重要提示:MMRotate默认采用OpenCV表示法。你在准备自定义数据集时,必须确保你的标注格式与之一致,或者在数据加载管道中进行正确的转换。很多开源遥感数据集(如DOTA)本身就采用这种格式。

2.2 构建自定义数据集的正确姿势

假设你有一批自己的航拍图或工业检测图,并用LabelImg、CVAT等工具标注了旋转框。你得到的标注文件可能是XML、JSON或TXT格式。你需要将它们转换为MMRotate能够读取的格式。

MMRotate主要支持两种数据格式:

  • DOTA格式:遥感领域权威数据集DOTA使用的格式。每张图片对应一个TXT标注文件。文件内容如下:

    x1 y1 x2 y2 x3 y3 x4 y4 category difficult

    其中(x1, y1), ..., (x4, y4)是旋转框四个顶点的坐标,按顺时针或逆时针顺序排列。category是类别名,difficult表示是否为难例(0或1)。MMRotate在读取时会自动将这四个顶点转换为默认的(cx, cy, w, h, angle)格式。

  • COCO格式:扩展了经典COCO格式以支持旋转框。在标准的annotations字段中,每个物体的bbox不再是一个4维数组[x, y, w, h],而是一个5维数组[x, y, w, h, angle]。这就是我们上面讨论的OpenCV表示法。segmentation字段可以存放多边形顶点,用于更精细的分割任务。

对于自定义数据集,我强烈推荐使用COCO格式。原因有三:1) 工具生态丰富,很多标注工具(如label-studio)可以直接导出或通过脚本转为COCO格式;2) 数据结构清晰,一个JSON文件管理所有标注,易于维护和版本控制;3) MMRotate对COCO格式的支持非常完善,数据加载效率高。

实操步骤:将你的标注转为COCO格式

假设你的原始标注是每张图一个XML(类似Pascal VOC),下面是一个大致的转换逻辑,你需要根据实际情况调整:

import json import xml.etree.ElementTree as ET import os import cv2 import numpy as np def voc_xml_to_coco_rotated(xml_dir, img_dir, output_json_path): """ 将VOC格式的旋转框XML转换为MMRotate可用的COCO格式。 假设XML中物体的bbox表示为: <bndbox><x>...</x><y>...<angle>...</angle>...</bndbox> 或者通过四个顶点<polygon>表示。 这里以四个顶点为例进行转换。 """ coco_dict = { "images": [], "annotations": [], "categories": [{"id": 1, "name": "your_class_name"}] # 根据你的类别修改 } image_id = 1 ann_id = 1 category_name_to_id = {"your_class_name": 1} for xml_file in os.listdir(xml_dir): if not xml_file.endswith('.xml'): continue tree = ET.parse(os.path.join(xml_dir, xml_file)) root = tree.getroot() # 处理image信息 filename = root.find('filename').text img_path = os.path.join(img_dir, filename) img = cv2.imread(img_path) if img is None: continue height, width = img.shape[:2] image_info = { "id": image_id, "file_name": filename, "height": height, "width": width, } coco_dict["images"].append(image_info) # 处理每个object for obj in root.findall('object'): class_name = obj.find('name').text # 假设旋转框由四个顶点表示 <polygon><x1>..<y1>..<x2>..<y2>..<x3>..<y3>..<x4>..<y4> polygon = obj.find('polygon') if polygon is not None: # 提取四个点坐标,注意顺序 points = [] for i in range(1, 5): x = float(polygon.find(f'x{i}').text) y = float(polygon.find(f'y{i}').text) points.append([x, y]) # 将四个顶点转换为 (cx, cy, w, h, angle) # 使用OpenCV的minAreaRect函数,它返回的正是 (cx, cy), (w, h), angle # 注意:cv2.minAreaRect 返回的angle范围可能是 [-90, 0),需确认与MMRotate默认一致 rect = cv2.minAreaRect(np.array(points, dtype=np.float32)) (cx, cy), (w, h), angle = rect # cv2.minAreaRect 返回的angle可能不符合MMRotate的(-90,0]约定,有时需要转换 # MMRotate中,通常使用 `cv2.minAreaRect` 后,角度已经基本符合,但为了保险,可以标准化 # 一个常见的转换函数: def norm_angle(angle, angle_range): """将角度标准化到指定范围。""" if angle_range == 'oc': # OpenCV表示法,限制在 (-90, 0] return angle if -90 < angle <= 0 else angle - 90 # 其他表示法... return angle angle = norm_angle(angle, 'oc') # 确保 w >= h,这是OpenCV表示法的隐含要求 if w < h: w, h = h, w angle -= 90 # 角度再次标准化 angle = norm_angle(angle, 'oc') bbox = [cx, cy, w, h, angle] else: # 如果没有多边形,可能是传统的水平框,需要你根据业务决定是否转为旋转框 # 通常,自定义数据集应该直接标注旋转框 continue ann = { "id": ann_id, "image_id": image_id, "category_id": category_name_to_id[class_name], "bbox": bbox, # 这里是5个数的列表! "area": w * h, "iscrowd": 0, "segmentation": [np.array(points).flatten().tolist()] # 可选,保存多边形 } coco_dict["annotations"].append(ann) ann_id += 1 image_id += 1 with open(output_json_path, 'w') as f: json.dump(coco_dict, f) print(f"转换完成,保存至 {output_json_path}") # 调用函数 voc_xml_to_coco_rotated('path/to/your/xmls', 'path/to/your/images', 'train.json')

这个脚本提供了一个转换框架。最关键的是角度标准化。你必须清楚你的原始标注工具输出角度的约定,并正确转换到MMRotate的OpenCV表示法。在实际操作中,建议先用几组数据可视化验证一下转换后的框是否与图像上的目标完美贴合。

3. 环境配置与MMRotate项目结构解析

工欲善其事,必先利其器。配置好环境并理解项目结构,能让你在后续的调试中事半功倍。

3.1 一步步搭建MMRotate环境

MMRotate依赖PyTorch、MMCV和MMDetection。下面是一个在Ubuntu 20.04/22.04上从零开始的稳定安装流程。使用Conda管理环境能避免很多依赖冲突。

# 1. 创建并激活conda环境(假设已安装Miniconda/Anaconda) conda create -n mmrotate python=3.8 -y conda activate mmrotate # 2. 安装PyTorch(请根据你的CUDA版本到PyTorch官网获取最新命令) # 例如,对于CUDA 11.3 conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch # 3. 安装MMCV(这是MM系列的基础库,必须版本匹配) # 使用mim安装是官方推荐的方式,它能自动处理依赖 pip install openmim mim install mmcv-full==1.7.1 # 版本需要与后续的MMDetection、MMRotate匹配 # 4. 安装MMDetection # 克隆MMDetection仓库(也可以直接pip install,但克隆方便查看源码和配置文件) git clone https://github.com/open-mmlab/mmdetection.git cd mmdetection # 切换到与MMRotate兼容的版本分支,MMRotate v1.0.0rc1 通常对应 MMDetection v2.25.0 git checkout v2.25.0 pip install -v -e . # “-e”表示以可编辑模式安装,方便修改源码 cd .. # 5. 安装MMRotate git clone https://github.com/open-mmlab/mmrotate.git cd mmrotate # 查看最新的发布版本或稳定分支 pip install -v -e . # 同样以可编辑模式安装 # 6. 验证安装 python -c "import mmrotate; print(mmrotate.__version__)" # 如果没有报错,并输出版本号,说明安装成功

避坑提示:版本兼容性是最大的坑!MMRotate、MMDetection、MMCV、PyTorch之间有着严格的版本对应关系。务必查阅你下载的MMRotate版本(如main分支或某个tag)的requirements.txt或官方安装文档,严格按照指定的版本号安装。盲目使用最新版大概率会失败。

3.2 项目目录与配置文件解读

安装完成后,MMRotate的目录结构如下:

mmrotate/ ├── configs/ # 所有模型的配置文件 │ ├── rotated_faster_rcnn/ │ ├── rotated_retinanet/ │ └── ... ├── mmrotate/ # 核心源代码 ├── tools/ # 训练、测试、推理脚本 ├── data/ # 建议在此处软链接或放置你的数据集 └── demo/ # 演示脚本

核心中的核心是configs/目录。MMRotate采用模块化的配置系统,一个完整的训练配置由多个基础配置文件组合而成。例如,configs/rotated_retinanet/rotated_retinanet_obb_r50_fpn_1x_dota_le90.py这个配置文件,通常会通过_base_字段继承多个基础配置:

_base_ = [ '../_base_/datasets/dota.py', # 数据集定义 '../_base_/schedules/schedule_1x.py', # 训练策略(学习率、迭代次数) '../_base_/default_runtime.py', # 运行时配置(日志、钩子) './rotated_retinanet_obb_r50_fpn.py' # 模型结构定义 ]

当你使用自定义数据集时,你需要重点关注并修改的就是数据集定义部分。你需要创建一个新的配置文件(例如my_custom_config.py),在其中通过_base_继承模型结构、训练策略等,然后重写(override)数据相关的配置。

为什么采用这种设计?这种设计将数据、模型、训练策略解耦,使得你可以像搭积木一样组合不同的实验设置。例如,你可以轻松地将为DOTA数据集设计的Rotated RetinaNet模型,通过只修改数据配置部分,应用到你的自定义数据集上,而无需改动模型代码。

4. 为自定义数据集创建配置文件

这是将你的数据“喂”给MMRotate的关键一步。我们假设你已经按照COCO格式准备好了数据,目录结构如下:

data/my_custom/ ├── train/ │ └── images/ # 存放所有训练图片 ├── val/ │ └── images/ # 存放所有验证图片 ├── train.json # 训练集标注文件 └── val.json # 验证集标注文件

接下来,在mmrotate/configs/下创建一个新目录,比如my_project/,然后新建配置文件rotated_retinanet_my_custom.py

# mmrotate/configs/my_project/rotated_retinanet_my_custom.py # 1. 继承基础配置 _base_ = [ '../rotated_retinanet/rotated_retinanet_obb_r50_fpn_1x_dota_le90.py', # 继承模型和基础设置 ] # 2. 修改数据集相关配置 dataset_type = 'MyCustomDataset' # 我们稍后会注册这个数据集类型 data_root = 'data/my_custom/' # 数据根目录 # 修改训练集配置 train_dataloader = dict( batch_size=2, # 根据你的GPU内存调整,旋转检测通常比水平检测更耗显存 num_workers=4, # 数据加载线程数,根据CPU核心数调整 persistent_workers=True, # 保持worker进程活跃,加速epoch间的数据加载 sampler=dict(type='DefaultSampler', shuffle=True), dataset=dict( type=dataset_type, data_root=data_root, ann_file='train.json', # 相对于data_root的路径 data_prefix=dict(img_path='train/images/'), filter_cfg=dict(filter_empty_gt=True), # 过滤掉没有标注的图片 # 非常重要:指定数据集的元信息,特别是类别名 metainfo=dict(classes=('building', 'vehicle', 'ship')), # 替换为你的类别名 pipeline=_base_.train_pipeline # 继承默认的训练数据增强流程 ) ) # 修改验证集配置 val_dataloader = dict( batch_size=1, num_workers=2, persistent_workers=True, sampler=dict(type='DefaultSampler', shuffle=False), dataset=dict( type=dataset_type, data_root=data_root, ann_file='val.json', data_prefix=dict(img_path='val/images/'), metainfo=dict(classes=('building', 'vehicle', 'ship')), test_mode=True, # 验证/测试模式,通常关闭数据增强 pipeline=_base_.test_pipeline ) ) # 修改测试集配置(通常和验证集一样) test_dataloader = val_dataloader # 修改评估器 val_evaluator = dict( type='DOTAMetric', # 或者 'CocoMetric',取决于你的评估需求 metric='mAP', # 主要评估指标为平均精度均值 iou_thrs=[0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95], # DOTA常用,可调整 # 如果你使用COCO格式,并且想用COCO的评估方式: # type='CocoMetric', # metric=['bbox', 'segm'], # 评估检测和分割(如果有) # ann_file=data_root + 'val.json', # classwise=True, # 输出每个类别的AP ) test_evaluator = val_evaluator # 3. 修改模型头中的类别数 # 找到模型配置中分类头和回归头的 `num_classes` 参数并修改 # 通常需要在 `model` 字典中覆盖 `bbox_head` 的 `num_classes` model = dict( bbox_head=dict( num_classes=3, # 修改为你的实际类别数,例如3类 ) ) # 4. 调整训练策略(可选) # 例如,根据你的数据集大小调整学习率和训练周期 # 默认的 `schedule_1x.py` 是针对12个epoch(在DOTA上)设计的。 # 如果你的图片数量远少于DOTA,可能需要减少总迭代次数或调整学习率衰减策略。 # _base_ 中的 schedule 配置可能被覆盖,这里直接修改优化器和学习率调度器 # optimizer optim_wrapper = dict( optimizer=dict(type='SGD', lr=0.005, momentum=0.9, weight_decay=0.0001)) # 初始学习率 # learning rate scheduler param_scheduler = [ dict( type='LinearLR', start_factor=0.001, by_epoch=False, begin=0, end=500), # 热身 dict( type='MultiStepLR', begin=0, end=12, # 总epoch数 by_epoch=True, milestones=[8, 11], # 在第8和第11个epoch降低学习率 gamma=0.1) ] # 5. 修改默认运行时配置(可选) default_hooks = dict( logger=dict(type='LoggerHook', interval=50), # 每50个iteration打印一次日志 checkpoint=dict(type='CheckpointHook', interval=1), # 每个epoch保存一次权重 ) # 训练总epoch数 train_cfg = dict(max_epochs=12, val_interval=1) # 每1个epoch验证一次

关键点解析与避坑

  1. metainfo中的classes:这里的类别顺序必须与你的标注文件(如COCO JSON)中categoriesid顺序一致。通常,在转换标注时,我们会让id从1开始连续编号,并与此处的classes元组一一对应。如果类别ID不连续或顺序错乱,会导致模型学习到错误的类别映射。

  2. 数据集类型MyCustomDataset:我们直接使用了MyCustomDataset,但MMRotate默认并没有这个类。我们需要注册这个数据集。最简单的方式是继承现有的支持COCO格式的数据集,如CocoDataset。在MMRotate中,通常使用RotatedDataset作为基类。我们需要在运行训练前,通过一段代码注册我们的数据集。

    创建一个文件mmrotate/datasets/my_custom.py

    from mmrotate.datasets import RotatedCocoDataset from mmrotate.registry import DATASETS @DATASETS.register_module() class MyCustomDataset(RotatedCocoDataset): """My custom dataset for rotated object detection.""" METAINFO = { 'classes': ('building', 'vehicle', 'ship'), # 这里定义类别名 'palette': [(220, 20, 60), (119, 11, 32), (0, 0, 142)] # 每个类别对应的颜色(可视化用) }

    然后,在你的训练脚本或配置文件中,需要先导入这个类,以确保它被注册到MMRotate的注册器中。一种更简单的方法是在配置文件中直接使用RotatedCocoDataset,并在metainfo中覆盖classes,就像我们上面做的那样。但为了清晰和可维护性,自定义一个数据集类是更好的实践。

  3. 评估器选择DOTAMetric是遥感领域常用的评估方式,它使用不同IoU阈值下的平均精度(mAP)。CocoMetric则是更通用的评估器。选择哪个取决于你的任务需求和领域习惯。如果你的数据集类似DOTA(目标方向多样、密集),使用DOTAMetric更合适。

5. 启动训练与监控

配置文件和数据准备就绪后,就可以开始训练了。MMRotate提供了强大的命令行工具。

5.1 单GPU训练

cd /path/to/mmrotate python tools/train.py configs/my_project/rotated_retinanet_my_custom.py \ --work-dir work_dirs/my_custom_exp \ # 指定工作目录,用于保存日志和模型 --cfg-options randomness.seed=42 \ # 设置随机种子保证可复现性 --auto-scale-lr # 根据实际batch_size自动缩放学习率(如果你修改了batch_size)

5.2 多GPU分布式训练

./tools/dist_train.sh \ configs/my_project/rotated_retinanet_my_custom.py \ 8 \ # 使用8个GPU --work-dir work_dirs/my_custom_exp_dist \ --cfg-options randomness.seed=42

5.3 训练过程监控

训练开始后,你可以在work_dirs/my_custom_exp目录下找到:

  • {timestamp}.log:详细的训练日志。
  • vis_data/:可视化数据,可用于TensorBoard。
  • {epoch}.pth:每个epoch保存的模型权重。

使用TensorBoard监控

tensorboard --logdir work_dirs/my_custom_exp --port 6006

然后在浏览器打开localhost:6006,你可以看到损失曲线、学习率变化、验证集mAP等关键指标,这对于调试和调参至关重要。

5.4 常见训练问题与调优

  1. Loss为NaN或突然爆炸

    • 原因:学习率过高、数据中存在异常标注(如坐标超出图像范围、宽高为负数或零)、梯度爆炸。
    • 排查
      • 首先大幅降低学习率(如lr=0.0001)试试。
      • 使用tools/analysis_tools/browse_dataset.py脚本可视化你的数据集,检查标注框是否正常。
      • 在配置文件中加入梯度裁剪:optim_wrapper = dict(clip_grad=dict(max_norm=35, norm_type=2))
  2. mAP始终为0或极低

    • 原因:类别定义不匹配、标注格式错误(特别是角度)、数据增强过于激进导致模型无法学习、模型复杂度与数据量不匹配(数据太少,模型太大)。
    • 排查
      • 确认标注:用可视化脚本检查转换后的(cx, cy, w, h, angle)是否正确。重点检查角度!一个快速验证方法是,在验证集上运行测试脚本,看看预测框是否大致在目标位置,哪怕类别错了。如果框都乱飞,肯定是定位出了问题,大概率是角度或中心点坐标的转换有误。
      • 简化流程:关闭所有数据增强(在配置文件中将train_pipeline替换为test_pipeline),在小批量数据上过拟合。如果模型能在几十张图上达到接近100%的训练精度,说明模型和数据管道基本是通的,问题可能出在增强或超参上。
      • 调整模型:对于小数据集,使用更小的 backbone(如r18代替r50)或更简单的检测头(单阶段如RetinaNet比两阶段Faster R-CNN更易训练)。
  3. 训练速度慢

    • 原因:图片分辨率过大、num_workers设置过小、使用了过大的模型。
    • 优化
      • 在数据管道中调整Resize步骤,将图片缩放到一个合理的尺寸(如scale=(1024, 1024))。注意保持长宽比,避免目标变形。
      • 适当增加num_workers(通常设置为CPU核心数的2-4倍),并启用persistent_workers=True
      • 考虑使用混合精度训练(AMP)。在配置文件中添加:optim_wrapper = dict(type='AmpOptimWrapper')。这能显著减少显存占用并加快训练速度,但对某些模型可能引入精度损失或不稳定,需要测试。

6. 模型测试、推理与可视化

训练完成后,你需要评估模型在独立测试集上的性能,并学会如何使用它进行预测。

6.1 测试模型性能

# 单GPU测试 python tools/test.py \ configs/my_project/rotated_retinanet_my_custom.py \ work_dirs/my_custom_exp/epoch_12.pth \ # 指定训练好的权重 --work-dir work_dirs/my_custom_exp/test_results \ # 保存测试结果 --show-dir work_dirs/my_custom_exp/test_vis \ # 保存可视化图片 --show # 同时显示图片(如果有GUI) # 多GPU测试 ./tools/dist_test.sh \ configs/my_project/rotated_retinanet_my_custom.py \ work_dirs/my_custom_exp/epoch_12.pth \ 8 \ --work-dir work_dirs/my_custom_exp/test_results_dist

测试脚本会加载模型和权重,在测试集上运行推理,并调用之前配置的test_evaluator(如DOTAMetric)计算mAP等指标。结果会保存在work_dirs/my_custom_exp/test_results下的JSON文件中,同时终端会打印出详细的评估报告。

6.2 对单张图片或视频进行推理

MMRotate提供了方便的演示脚本。你可以编写一个简单的Python脚本:

from mmrotate.apis import inference_detector, init_detector import cv2 import numpy as np # 1. 初始化模型 config_file = 'configs/my_project/rotated_retinanet_my_custom.py' checkpoint_file = 'work_dirs/my_custom_exp/epoch_12.pth' model = init_detector(config_file, checkpoint_file, device='cuda:0') # 或 'cpu' # 2. 读取图片 img = 'test_image.jpg' image = cv2.imread(img) # 3. 进行推理 result = inference_detector(model, image) # 4. 可视化结果 # 你可以使用MMRotate内置的可视化函数 from mmrotate.visualization import RotLocalVisualizer visualizer = RotLocalVisualizer() visualizer.set_image(image) # 假设result是一个包含预测框和分数的列表 # 需要根据模型输出结构进行解析,这里是一个示例 pred_instances = result.pred_instances bboxes = pred_instances.bboxes.cpu().numpy() # 旋转框 [N, 5] labels = pred_instances.labels.cpu().numpy() # 类别 [N] scores = pred_instances.scores.cpu().numpy() # 分数 [N] # 过滤低分预测 score_thr = 0.3 keep = scores > score_thr bboxes = bboxes[keep] labels = labels[keep] scores = scores[keep] # 绘制框 for i, bbox in enumerate(bboxes): # bbox: [cx, cy, w, h, angle] # 将旋转矩形转换为四个顶点,用于绘制 from mmrotate.structures.bbox import rbox2qbox qbox = rbox2qbox(bbox) # 转换为四个顶点 [x1,y1,x2,y2,x3,y3,x4,y4] pts = qbox.reshape(4, 2).astype(np.int32) cv2.polylines(image, [pts], isClosed=True, color=(0, 255, 0), thickness=2) # 添加标签和分数 label_text = f'{model.dataset_meta["classes"][labels[i]]}: {scores[i]:.2f}' cv2.putText(image, label_text, (int(bbox[0]), int(bbox[1])-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) # 5. 保存或显示结果 cv2.imwrite('result.jpg', image) # cv2.imshow('Result', image) # cv2.waitKey(0)

6.3 结果分析与模型优化

得到测试结果后,不要只看一个总的mAP。深入分析对于改进模型至关重要。

  1. 逐类别AP分析:查看每个类别的AP。如果某个类别(如“ship”)的AP远低于其他类别,可能的原因有:该类别样本数量太少(类别不平衡)、目标尺度特殊(太大或太小)、目标角度分布特殊。解决方案包括:数据增强(针对小目标)、收集更多该类别数据、使用Focal Loss缓解类别不平衡、为该类别设置独立的检测头参数。

  2. 可视化错误案例:将测试集上预测错误的图片(漏检、误检、定位不准)可视化出来。这是发现模型短板的直接方法。

    • 漏检:目标太小?被遮挡?与背景颜色/纹理相似?考虑在数据增强中加入更多针对性的变换(如随机裁剪、Mosaic增强),或者调整模型Anchor的尺度。
    • 误检:背景区域被误认为目标?可能是负样本不足或特征混淆。可以尝试在训练时加入困难负样本挖掘(OHEM),或者调整分类损失函数的权重。
    • 定位不准:框的角度或中心点偏移。旋转框回归比水平框更复杂。可以检查回归损失函数(通常是Smooth L1 Loss)是否收敛正常。对于角度回归,有些研究使用离散化(将角度分为多个bin)或使用IoU-Smooth L1损失来提升精度。
  3. 超参数调优:基于分析结果,可以系统地调整超参数。MMRotate支持方便的配置覆盖。例如,你想尝试不同的Anchor尺寸:

    python tools/train.py configs/my_project/rotated_retinanet_my_custom.py \ --work-dir work_dirs/exp_tune_anchor \ --cfg-options model.bbox_head.anchor_generator.scales="[4, 8, 16, 32, 64]" \ model.bbox_head.anchor_generator.ratios="[0.5, 1.0, 2.0]" \ model.bbox_head.anchor_generator.rotations="[-90, -60, -30, 0]"

    通过--cfg-options,你可以覆盖配置文件中任何深层次的参数,无需修改原文件,非常适合进行实验对比。

7. 进阶技巧与部署考量

当你的基础模型跑通后,可以考虑以下进阶优化,并思考如何将模型投入实际应用。

7.1 数据增强策略优化

MMRotate集成了MMDetection丰富的数据增强管道。对于旋转目标检测,一些特定的增强非常有效:

  • RandomRotate:这是旋转检测的“本命”增强。它能随机旋转图片和对应的旋转框,极大地增加模型对目标方向的鲁棒性。在配置文件的train_pipeline中加入:

    dict(type='RandomRotate', prob=1.0, level=10, max_rot_angle=180),

    注意,max_rot_angle设置过大可能导致目标旋转后与标注的角度表示范围冲突,需要谨慎。

  • MosaicMixUp:这些混合多张图片的增强技术能有效提升模型对小目标和遮挡目标的检测能力,尤其适用于数据量不大的场景。

  • PhotoMetricDistortion:随机调整亮度、对比度、饱和度、色调,模拟不同的光照条件,提升模型泛化性。

一个强化后的train_pipeline可能看起来像这样(在配置文件中修改):

train_pipeline = [ dict(type='LoadImageFromFile'), dict(type='LoadAnnotations', with_bbox=True, box_type='qbox'), # 注意box_type dict(type='Resize', scale=(1024, 1024), keep_ratio=True), dict(type='RandomFlip', prob=0.5), dict(type='RandomRotate', prob=1.0, level=10, max_rot_angle=60), # 加入随机旋转 dict(type='PhotoMetricDistortion'), dict(type='PackDetInputs') ]

7.2 模型选择与集成

MMRotate提供了多种旋转检测算法:

  • Rotated RetinaNet:单阶段,速度快,适合实时性要求高的场景。
  • Rotated Faster R-CNN:两阶段,精度通常更高,但速度较慢。
  • RoI Transformer:专门为旋转目标设计,在RPN(区域提议网络)阶段就生成旋转的提议框,精度提升明显,尤其适合密集、小目标场景。
  • Oriented R-CNN:另一种高效的两阶段旋转检测器,设计更简洁。

对于你的自定义数据集,建议先用Rotated RetinaNet快速验证流程和基线性能。如果精度不满足要求,再尝试RoI TransformerOriented R-CNN。可以在相同的数据和训练设置下比较它们的精度-速度曲线,选择最适合的模型。

7.3 模型部署:从PyTorch到生产环境

训练好的PyTorch模型(.pth文件)不能直接用于生产环境(如C++服务、移动端、边缘设备)。你需要将其转换为推理引擎支持的格式。

主流部署路径:

  1. ONNX -> TensorRT (NVIDIA GPU):这是NVIDIA平台上的黄金组合。

    • 步骤: a. 使用MMDeploy(OpenMMLab的部署工具箱)或PyTorch原生torch.onnx.export将模型转换为ONNX格式。注意:转换旋转检测模型时,需要确保模型中的自定义算子(如旋转NMS)也被正确导出和支持。 b. 使用TensorRT的trtexec工具或Python API将ONNX模型编译为高度优化的TensorRT引擎(.engine文件)。
    • 优势:极致性能,低延迟。
    • 挑战:算子支持,特别是后处理中的旋转NMS,可能需要自定义插件(Custom Plugin)。
  2. TorchScript:PyTorch自带的序列化格式,可以在没有Python环境的C++中运行。

    • 步骤:使用torch.jit.tracetorch.jit.script将模型转换为TorchScript。
    • 优势:流程相对简单,与PyTorch生态结合紧密。
    • 挑战:某些动态控制流或复杂数据结构的模型可能无法完美转换。性能优化不如TensorRT。
  3. OpenVINO (Intel CPU/GPU):针对Intel硬件优化的部署工具包。

    • 步骤:将ONNX模型通过OpenVINO的Model Optimizer转换为IR格式,再用Runtime推理。
    • 优势:在Intel CPU上性能优秀,易于使用。

部署建议

  • 前期验证:先使用MMRotate自带的inference_detectorAPI搭建一个简单的Python服务(如Flask/FastAPI),快速验证模型功能和性能。
  • 性能瓶颈分析:使用 profiling 工具(如PyTorch Profiler, NVIDIA Nsight)分析推理过程中各阶段耗时。通常,后处理(尤其是NMS)在旋转检测中占比不低。
  • 自定义后处理:生产环境中,为了极致性能和控制力,经常需要将模型的后处理(解码bbox、NMS)从模型中剥离,用高度优化的C++代码实现。你需要理解模型原始输出的格式(如每个Anchor的类别得分和bbox偏移量),并重新实现解码和NMS逻辑。

7.4 持续集成与模型管理

当项目迭代起来后,模型版本、配置文件、训练日志会越来越多。建议建立简单的管理规范:

  • 版本控制:使用Git管理你的配置文件、数据集转换脚本和自定义代码。
  • 实验跟踪:使用MLflow、Weights & Biases或TensorBoard来记录每次实验的超参数、指标和模型文件。
  • 自动化:编写脚本自动化训练-评估-测试的流程,特别是当你要进行大规模超参数搜索时。

从准备自定义数据到训练、调优、最终部署一个旋转目标检测模型,整个过程充满了挑战,但也正是这些挑战让结果更有价值。MMRotate作为一个强大的工具箱,为你扫清了许多底层实现的障碍,让你能更专注于数据、模型结构和业务逻辑本身。希望这篇基于实战的详细指南,能帮助你顺利跨过旋转目标检测的门槛,在实际项目中取得成功。记住,遇到问题时,多可视化、多分析数据、从小规模实验开始,是最高效的调试路径。