RTDETR_OBB利用transform直接预测旋转框(Ultralytics_yolov11版本)

一、RTDETR(u版)流程回顾

一文带你搞定Ultralytics(yolov11版本)中的RTDETR训练以及预测的代码解析和流程分析

二、RTDETR_OBB(u版)细节

1、RTDETRDecoder初始化

#encoder和decoder head部分预测bbox的部分需要修改成直接预测rbbox
self.dec_bbox_head = nn.ModuleList([MLP(hd, hd, 4, num_layers=3) for _ in range(ndl)])
self.dec_bbox_head_obb = nn.ModuleList([MLP(hd, hd, 5, num_layers=3) for _ in range(ndl)])

self.enc_bbox_head = MLP(hd, hd, 4, num_layers=3)
self.enc_bbox_head_obb = MLP(hd, hd, 5, num_layers=3)

2、构建obb去噪组get_cdn_group_obb

下图是rtdetr预测矩形框的去噪组,去噪组分为正样本和负样本。
正样本在gt的基础上加上了原gt的0-h/2和0-w/2的扰动。负样本加上了原gt的h/2-h和w/2-h的扰动。具体值也可调参。
在这里插入图片描述
旋转框只需要在gt的wh基础上额外加上一个角度theta的扰动即可。

3、替换loss

所有计算bbox的loss替换成计算rbbox的probiou。(ps:匈牙利匹配里的计算bbox_iou也要替换成probiou来计算rbbox_iou)

if self.use_obb:
    loss[name_giou] = 1.0 - probiou(pred_bboxes, gt_bboxes)
else:
    loss[name_giou] = 1.0 - bbox_iou(pred_bboxes, gt_bboxes, xywh=True, GIoU=True)
### 实现 RTDETR 模型中的 OBB 支持 为了使 RTDETR 模型能够处理定向边界 (OBB),需要对现有模型架构进行一系列修改。这些改动主要集中在数据预处理、损失函数调整以及预测器模块的设计上。 #### 数据预处理阶段 对于 OBB 的支持,输入图像的目标标注不再是简单的矩形坐标(xmin, ymin, xmax, ymax),而是包含角度信息的五参数表示形式(cx, cy, w, h, θ)[^1]。因此,在数据加载部分需增加转换逻辑: ```python def convert_rect_to_obb(rect): """Convert standard rectangle to oriented bounding box.""" cx = (rect[0] + rect[2]) / 2.0 cy = (rect[1] + rect[3]) / 2.0 width = abs(rect[2] - rect[0]) height = abs(rect[3] - rect[1]) # Assuming the angle is initially set as zero; this will be adjusted during training. theta = 0 return [cx, cy, width, height, theta] ``` #### 修改损失计算方法 传统的 IoU 计算方式不再适用于旋转后的边情况。为此引入 RotatedIoULoss 或者其他更适合于评估两个任意方向上的矩形重叠程度的方法来替代原有的 L1 Loss 和 GIoU Loss: ```python from mmdet.models.losses import build_loss class RotatedIoULoss(nn.Module): def __init__(self, **kwargs): super(RotatedIoULoss, self).__init__() def forward(self, pred, target, weight=None, avg_factor=None, reduction_override=None): ... # Implementation of rotated IOU loss calculation here ... ``` #### 自定义预测器组件 参照 Ultralytics YOLO 中 `OBBPredictor` 类的做法,在 RTDETR 上创建类似的子类以覆盖默认的行为并加入特定于 OBB 预测的功能: ```python from rt detr.predictors.detection_predictor import DetectionPredictor class OBBDetectionPredictor(DetectionPredictor): def post_process(self, preds_dict): """ Post-process predictions into objects with orientation information. Args: preds_dict(dict): Dictionary containing raw prediction tensors Returns: list[List[np.ndarray]]: A nested list where each sublist contains numpy arrays representing detected instances, including their class labels and obb parameters. """ results = [] for i in range(len(preds_dict['boxes'])): boxes = preds_dict['boxes'][i].detach().cpu().numpy() scores = preds_dict['scores'][i].detach().cpu().numpy() classes = preds_dict['labels'][i].detach().cpu().numpy() dets = np.hstack((boxes[:, :5], scores[:, None])) # Add score column after five bbox params keep_idx = nms_rotated(dets, cfg.test_cfg.nms.iou_thr) result = dict( bboxes=dets[keep_idx], labels=classes[keep_idx]) results.append(result) return results ``` 通过上述三个方面的改进,可以有效地让 RTDETR 模型具备识别和定位具有不同朝向物体的能力。
评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值