TensorFlow目标检测中带角度边界框的表示与预测方案咨询

刚好之前折腾过在TensorFlow Object Detection API里加旋转框检测的需求，给你整理几个可行的落地思路：

• 改造数据标注与输入管道
  首先得把标注体系从轴对齐框的xmin, ymin, xmax, ymax改成旋转框格式，推荐用cx, cy, w, h, theta（中心坐标、宽高、旋转角度，theta一般是框相对于水平轴的旋转角）。接下来要适配API的输入流程：

  • 重写tfrecord生成脚本，把旋转框的5个参数都写入样本的标注字段；
  • 修改输入读取的解析函数，在解析样本时额外把theta参数读出来，确保训练时模型能拿到完整的标注目标。

• 替换为旋转框专属的损失函数
  原API的损失都是针对轴对齐框设计的，比如普通IOU、Smooth L1，得换成适配旋转框的版本：

  • 核心是实现**旋转IOU（Rotated IOU）**损失，计算预测旋转框和真实框的交并比作为损失项，这是旋转检测里最常用的损失逻辑；
  • 角度参数的约束要注意周期性（比如0°和360°是同一个状态），可以把直接预测theta改成预测sinθ和cosθ，用Smooth L1损失约束这两个值，避免角度跳变带来的训练不稳定。

• 调整模型的输出头结构
  标准检测模型的输出层只输出4个框参数，你得把输出维度改成5（对应cx, cy, w, h, theta）或者6（sinθ+cosθ代替theta）：

  • 比如在SSD、Faster R-CNN这类模型里，找到负责预测框的卷积输出层，把通道数从原来的4类别数改成5类别数（或者6*类别数）；
  • 如果用CenterNet这类关键点检测框架，就在关键点的回归分支里加一个角度预测的子分支，输出旋转角度相关参数。

• 适配旋转框的后处理逻辑
  模型预测出旋转框参数后，后处理也得跟着改：

  • 把原来的非极大值抑制（NMS）换成旋转NMS，普通NMS没法处理旋转框的重叠判断；
  • 可视化模块也要调整，原来的cv2.rectangle只能画轴对齐框，得用cv2.drawContours先计算旋转矩形的四个顶点，再绘制轮廓。

• 参考成熟的旋转检测方案少走弯路
  不用完全从零造轮子，可以借鉴开源旋转检测算法的思路：

  • 参考DOTA数据集常用的旋转检测模型（比如RRPN、Rotated RetinaNet）的结构，在TF OD API的基础模型上做对应修改；
  • 有些开源项目已经做过类似的扩展，可以参考它们的损失函数、输出头实现逻辑。

内容的提问来源于stack exchange，提问作者Jenny