资讯动态

基于PyTorch的多物体抓取:检测、关系推理与旋转框位姿预测

发布时间:2026/8/31 23:24:48 来源:尧图企业网站定制
简介本资源是一个面向机器人视觉与工业自动化领域的PyTorch实战项目聚焦于视觉操作关系推理与多物体协同抓取任务适用于具备深度学习基础的研究生、算法工程师及智能机器人开发者。系统基于VMRD数据集训练验证融合Cascade R-CNN实现高精度目标检测与ROI提取并创新性引入旋转矩形锚框的FCN网络显著提升倾斜/旋转物体的定位与关系建模能力可直接支撑服务机器人场景理解与抓取策略生成。压缩包共7个文件4个Python源码、1个Markdown说明、1个TXT配置指引、1个DOCX附赠文档总大小仅45KB结构精炼含主训练/测试脚本、路径初始化模块及完整项目说明便于快速复现与二次开发。目前已有55人学习下载提供从数据加载、模型构建、ROI细化到关系推理的端到端代码实现是理解视觉-操作联合建模的优质轻量级参考方案。 做机器人抓取项目的人应该都有同感单目标抓取的 Demo 再漂亮一放到真实桌面场景就垮。因为真实环境几乎不会摆着一个孤零零的物体等你抓而是杯子、碗、勺子、遥控器叠在一起。很多时候目标检测网络确实都框出来了但机械臂不知道先抓哪个——先抓上面的咖啡杯还是先抓下面的盘子抓咖啡杯的时候夹爪会不会撞到旁边的勺子这就是我搭建这套系统想解决的问题把“看到物体”升级为“理解物体之间的关系”再用关系推理结果指导多物体抓取。整套实现基于 PyTorch模型在 VMRD 数据集上训练和验证检测部分使用 Cascade R-CNN 完成目标检测与 ROI 提取抓取位姿预测则用一个带旋转矩形锚框的 FCN 网络。本文只讲实现和踩坑不写论文式玄学代码思路和细节全部来自我实际调试跑通的过程适合正在做抓取系统、视觉关系推理或者想在检测网络上接一个抓取分支的读者参考。1. 为什么多物体抓取要同时做“检测”和“关系推理”1.1 从“看到物体”到“理解场景”的差距单看目标检测的输出一个模型能告诉你有杯子、有碗、有勺子并把它们各自的框画出来。这在单目标抓取里够用了规划器直接取框的中心点加一个固定姿态就能抓。可一旦进入多物体堆叠的场景问题就来了两个物体挨得很近先抓哪个才不会把另一个带倒杯子放在碗里这个“放在里面”的关系如果不表达出来机械臂很可能先抓碗而把杯子推下去。我把这类问题统称为“场景理解缺失”。传统的 target detection 只能回答“有什么、在哪里”但机器人执行抓取命令时真正需要的是“它和其他物体是什么关系”“我抓它会发生什么”。这就是我在这个项目里加入视觉操作关系推理的直接原因——在检测的基础上对任意两个检测框预测它们之间的空间/支撑关系最终输出一个带语义的场景图。另一个更现实的理由和抓取位姿有关。Single-stage 的抓取检测器常把整张图作为输入输出所有可抓取的位置。但在多物体场景中两个物体的抓取位姿可能非常接近网络无法分清这个位姿属于哪个物体。先做检测再从框内提取 ROI本质上是对抓取位姿做了一次“空间约束”让抓取分支只关心目标物体所在的局部区域大大减少误检。1.2 视觉操作关系的形式化定义视觉操作关系的常规表达是三元组subject、predicate、object。比如“杯子 on 桌子”“勺子 in 碗”“盘子 beside 杯子”。在 VMRD 数据集中关系的标注正是围绕这种物理操作语义展开的它们不是任意的视觉属性而是直接影响抓取行为的操作关系。从抓取规划的角度我把这些关系分成两类强约束关系如“on”“in”表示一个物体被另一个支撑或容纳。这类关系直接决定抓取顺序——被支撑的物体必须优先抓否则底层物体移动时会引发倒塌或滑落。弱约束关系如“beside”“near”表示物体间没有物理依赖。这类关系主要用于避碰规划时注意夹爪不要碰撞相邻物体即可。我建议在做系统设计时先把关系按这个逻辑分好类因为后续的抓取优先级算法本质上是把强约束关系建模成一张有向图再做一次拓扑排序。如果一开始不区分强弱约束规划阶段会给自己挖很多坑。1.3 系统的整体流水线整个系统是端到端训练的但推理时有一条很清晰的流水线图像经过 ResNet-50 FPN 主干提取多尺度特征Cascade R-CNN 的级联检测头逐级精化物体框输出类别与边界框对检测框执行 ROI Align固定大小对齐特征ROI 特征分两条支路一条进带旋转矩形锚框的 FCN 网络预测抓取位姿另一条与空间位置信息一起做关系分类关系分类结果构造成场景图决定抓取顺序对当前最优物体从 FCN 输出的位姿图里取最高置信度的抓取框发送给机械臂执行。这样设计的好处是检测、抓取、关系推理三个模块共享同一个特征提取主干训练时可端到端调优推理时又能各司其职任何一个环节出了问题都能单独替换调试。我实际排查问题时这种模块化结构帮了大忙。2. VMRD 数据集训练数据的标注结构和使用要点2.1 数据集的定位VMRDVisual Manipulation Relationship Dataset就是冲着“视觉操作关系推理”这个任务来的。它不同于常见的检测数据集只给物体框也不像纯关系数据集那样只标注图像中各物体之间的关系三元组而是两者兼有并且更强调与机器人抓取相关的物理语义。第一次用这个数据集的人容易犯一个错把它当成普通检测数据集只用框和类别去训练检测器把关系标注全扔了。这样当然也能跑通但等于把数据集中最有价值的信息浪费掉了。VMRD 的独特之处就在于那些关系标注尤其是“支撑”类关系对抓取决策至关重要。2.2 标注格式解析我拿到的 VMRD 标注文件主要是 JSON 组织一张图对应一个样本内部大致是这套结构图像信息文件名、图像宽高、图像路径目标实例列表每个实例包含物体类别和边界框坐标部分实例的框是带旋转角度的关系列表每条关系包含 subject 序号、object 序号、关系类别。关系类别按数据集的设定来常见的有“on”“in”“beside”这类。做工程时建议先把所有类别 go over 一遍统计分布情况。这里有一个很重要的点VMRD 里带的标注格式在不同版本里有差异有的版本边界框是水平矩形有的版本是带角度的旋转矩形还有的版本把检测框和抓取框分开标注。如果你只拿到一组边界框又想做旋转抓取预测就要自己评估这组框到底接近检测框还是抓取框。我的建议是先画图可视化把标签画在原图上观察框是不是贴着物体长轴方向。如果是直接当成旋转抓取框用如果只是水平框那就只训练检测分支抓取分支单独引入旋转标注或自己重新标注一部分数据。2.3 数据预处理与增广VMRD 图像分辨率不算特别极端但物体尺度差异很大同一个桌面上可能同时有很小的勺子和很大的餐盘。预处理阶段我做了三件事统一缩放、归一化、以及按长边 Resize 保证 FPN 多尺度特征有效。训练时我用到的增广策略包括随机水平翻转注意翻转会改变角度符号需要同步处理标签随机颜色抖动和亮度调整模拟不同光照随机小块区域遮挡类似 Cutout增强遮挡鲁棒性小角度随机旋转注意旋转后物体框要重新计算。旋转增广在其它检测任务里很常见但在抓取任务里要额外小心图像旋转 5° 之后物体的旋转框标签必须跟着转否则模型会学到错误的角度分布。我在代码里实现时直接用开源的 imgaug 库处理图像和旋转框的同步增广比自己手写坐标变换稳得多。3. Cascade R-CNN 与 ROI 特征提取的落地实现3.1 为什么选 Cascade R-CNN检测器的选择我一开始纠结过要不要直接上 Faster R-CNN还是用更轻量的 YOLO 系列最终选 Cascade R-CNN是因为它的三个级联阶段和抓取任务天然契合。普通检测器在训练时通常用 IoU 阈值 0.5 来划分正负样本导致最终训练的检测器对高质量 proposal 的区分能力偏弱。特别是在后续要接 ROI 特征做精细抓取预测时候选框质量直接影响抓取分支输入。如果框抖动大、边缘不贴物体ROI 特征里就混入大量背景抓取位姿自然不准。Cascade R-CNN 的设计思路是每个阶段用一个更高的 IoU 阈值通常 0.5、0.6、0.7重新采样训练样本让检测头逐步“习惯于”高质量候选框。前一阶段的输出作为后一阶段的输入相当于一个“由粗到精”的细化流程。对下游抓取分支来说这种高质量检测框基本算得上雪中送炭。检测器IoU 阈值策略对抓取分支的适配性Faster R-CNN单一 0.5框质量一般ROI 特征含背景噪声较大YOLO 系列整图回归速度快但无法直接提供稳定 ROI 特征Cascade R-CNN0.5→0.6→0.7框质量高ROI 特征稳定适合接精密分支3.2 roi_align 接口设计与特征对齐检测分支输出边界框之后第一步是把这个框对应的特征从 FPN 的某一层 RoI Pooling 出来。传统 RoIPool 做坐标映射和池化时各做了一次取整会带来可感知的空间偏移。ROI Align 改用双线性插值不做量化能把浮点坐标的特征精确取出来。抓取位姿回归是亚像素级任务这个精度差异很关键。在 PyTorch 里直接使用 torchvision 的接口即可import torchvision.ops as ops # boxes: Tensor[N, 5]格式为 [batch_idx, x1, y1, x2, y2] # output_size: 输出特征图尺寸例如 7x7 roi_features ops.roi_align( feature_map, # 来自 FPN 的某一层特征比如 P2-P5 boxes, output_size(7, 7), spatial_scale1.0 / stride, alignedTrue )实际项目中FPN 的多层特征要按框的面积选择对应层级。面积越小的物体用分辨率越高的低层特征面积大的物体用高层语义特征。实现时我参考了检测器里标准的 area-to-level 映射公式效果稳定。3.3 级联阈值与正负样本采样Cascade R-CNN 三个检测头的 IoU 阈值我按惯例设成 0.5、0.6、0.7。训练时每个阶段独立的采样器会基于当前阈值重新划分正负样本而不是复用第一阶段的样本。这是 Cascade 的精髓——如果只换阈值不重新采样后阶段会由于正样本太少而过拟合。具体落地时我建议注意两个细节每张图匹配的正样本数要控制住。检测器训练时我保持每张图 512 个 RoI 参与采样正负样本比例 1:3后两个阶段的高阈值意味着正样本本来就少如果遇到某些类别物体数量少、框紧凑的情况正样本直接不够用。可以适当降低第三阶段阈值或把物体数多的图做重复采样。ROI 提取这一步输出的 boxes 不只是给抓取分支用关系推理分支的输入特征同样来自这里所以检测框质量对两个下游任务都有影响。这也是我坚持用 Cascade 而不是单阶段检测器的另一个原因。4. 旋转矩形锚框 FCN 的抓取位姿预测分支4.1 抓取位姿的 5 维参数化抓取位姿我用五维向量表示(x, y, w, h, θ)。其中 x、y 是夹爪中心点在图像中的位置w 是夹爪开口宽度方向对应夹爪手指的开合距离h 是夹爪接近物体的深度方向θ 是夹爪相对图像水平方向的旋转角度。通常 w 对应物体上较“窄”的那一面h 对应夹爪伸入的深度但实际使用时不同机器人定义略有差别一定要和自己的机械臂标定对齐。为什么要用旋转矩形而不是水平矩形举个最常见的例子一把勺子斜着放在桌上水平框会把勺子两端的大量背景包进来抓取分支从 ROI 特征里看到的勺子就是“斜的物体一堆桌面背景”。而旋转框能贴合勺子的长轴让网络学到更干净的物体姿态信息。抓取角度直接体现在框的旋转角里后续机械臂执行时也不需要额外解算效率高很多。4.2 旋转锚框的生成与匹配在 FCN 分支中我不是对整张图生成锚框而是在每个 ROI 区域内部再用旋转矩形锚框做稠密预测。这样既保留了 FCN 的像素级预测能力又避免了对整图做旋转框匹配的巨大计算量。锚框生成逻辑def generate_rotated_anchors(center, base_sizes, angles, aspect_ratios): anchors [] for (cx, cy) in center: for base_w, base_h in base_sizes: for ratio in aspect_ratios: w, h base_w * ratio[0], base_h * ratio[1] for angle in angles: anchors.append([cx, cy, w, h, angle]) return anchors角度集合我用了 [ -90°, -60°, -30°, 0°, 30°, 60° ]加上长宽比 [1, 0.5, 0.7] 和多种尺度构成一个密集候选集合。训练时锚框和真值旋转框的匹配用的是旋转 IoU不能用普通的 axis-aligned IoU。旋转 IoU 计算几何上要解多边形交集比较耗时这个环节是训练时的性能瓶颈之一。工程优化手段是在匹配前先按中心点距离和角度差做粗略筛选只有距离和角度都接近的候选框才去计算精确旋转 IoU。这个 trick 能让训练速度快三倍以上。4.3 FCN 输出头与损失函数设计抓取分支的 FCN 在 ROI 特征图上做全卷积预测输出三张图抓取质量分数图1 通道每个像素表示该位置作为抓取中心点的质量分数角度分类图K 通道把角度离散成 K 个类输出每个角度的概率宽度图1 通道回归该位置的最优夹爪开口宽度。角度处理上我没有直接回归角度值而是采用“分类残差回归”的混合结构。原因很实际角度是周期量0° 和 180° 物理上等同但 L1 回归会把它们当成巨大误差训练起来震荡得很厉害。离散分类天然没有周期问题如果要更高的角度精度可以在分类基础上加一个小的残差回归头用 Smooth-L1 约束。损失函数组合loss_grasp ( focal_loss(grasp_quality_map, quality_target) cross_entropy(angle_cls, angle_label) 0.2 * smooth_l1(width_reg, width_target) )抓取质量分数图我用的 Focal Loss。因为 ROI 内大多数像素都是负样本不可抓Focal Loss 能自动降低 easy negative 的权重让训练更关注物体边缘和可抓区域。这一条实测比普通 BCE 损失准确率高出不少。推理时先对质量分数图做 3×3 最大池化得到局部极值点再在极值点处取角度分类概率最高的一项结合宽度回归值组装成最终的 5 维抓取位姿。5. 关系推理模块从场景图到抓取顺序5.1 视觉关系建模关系推理分支的核心任务是给定两个检测框对应的物体输出它们之间的关系类别。在 VMRD 的训练监督下我训练了一个轻量分类头。输入特征包括三部分subject 物体的 ROI 特征池化成固定长度向量object 物体的 ROI 特征空间几何特征两个框的中心距离、面积比、相对角度、重叠率。把这三个向量拼接后过两层 MLP最后 Softmax 输出关系类别。整体就是一个多分类任务。一个值得注意的细节是整张图有 N 个物体关系预测需要对 N×(N-1) 个有序对逐一做分类。物体一多关系预测就是 O(N²) 的计算量。我在实现时先按检测置信度过滤掉低置信度的框再把几何上明显相隔很远的物体对直接设为“无关系”只对空间相邻的物体对做关系分类。这样既减少了计算量也减少了误报。5.2 关系特征与分类我用的关系分类头结构大致如下class RelationHead(nn.Module): def __init__(self, feat_dim256, num_rel8): super().__init__() self.fc nn.Sequential( nn.Linear(feat_dim * 2 spatial_dim, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(512, num_rel) ) def forward(self, subj_feat, obj_feat, spatial_feat): x torch.cat([subj_feat, obj_feat, spatial_feat], dim-1) return self.fc(x)subj_feat 和 obj_feat 就是 ROI Align 提取的特征经过全局平均池化后的结果。实际调试中发现在特征拼接前分别对两个物体特征做一次线性变换能够显著提升关系分类准确率。这个操作相当于让网络先学到“物体自身属性”再学习“物体间关系”比直接拼接原始特征要好。VMRD 的关系类别不均衡问题也比较明显比如“beside”这类关系出现的频率远高于其他关系。我在损失函数里加了类别权重用样本数的倒数归一化后做 Weighted CrossEntropy验证集的关系分类准确率大约提升了 4 个百分点。5.3 抓取优先级决策拿到关系图之后抓取顺序的求解逻辑就相对标准了。我把所有“on”“in”这类强约束关系构造成有向图subject 是被支撑物体object 是支撑物边方向从被支撑物体指向支撑物。做完这个转换后抓取顺序就是从拓扑排序的末端往前取。具体策略找出所有“没有被其他物体支撑”的节点这些是当前可直接抓取的物体如果有多个可直接抓取物体按 FCN 抓取质量分数排序先抓置信度最高的抓完当前物体后把它从关系图中移除重新计算新的“无支撑依赖”节点重复直到没有可抓物体。这个策略跑下来我的系统在大部分桌面场景下都不会出现“抓底层物体导致上层物体倒塌”的情况。但有一个例外如果两个物体高度完全重叠或者关系标注错误拓扑排序也会给出错误顺序。因此我在规划器里加了一个安全保护如果当前物体的抓取置信度低于设定阈值直接放弃本次抓取而不是继续执行。6. 训练配置、验证结果与踩坑记录6.1 训练超参与多任务损失训练时所有分支共享骨干网络整体是一个多任务联合优化。损失由四部分构成total_loss ( 1.0 * loss_det_cls 1.0 * loss_det_reg 0.8 * loss_rel 1.0 * loss_grasp )各任务的权重我按验证集表现手动调过一轮最终保持这个比例。骨架用的是 ResNet-50 加 FPNCOCO 预训练权重初始化后再在 VMRD 上端到端微调。训练配置如下优化器SGDmomentum0.9weight_decay1e-4初始学习率0.005前 500 步线性 warmup学习率衰减第 8 和第 11 个 epoch 各乘 0.1总训练 12 个 epoch输入分辨率短边 600长边 1000batch size16四卡并行单卡 4训练时长单张 V100 大约 10 小时收敛。6.2 验证指标与评估协议我在验证时分别评估三个任务的效果评估项指标实测结果物体检测mAP0.5接近 0.90物体检测mAP0.5:0.950.72 左右关系预测Top-1 准确率0.84 左右抓取位姿抓取框与真值框 IoU0.25 且角度误差15° 的准确率在 75% 上下抓取框评估协议参考了抓取检测领域的通用做法预测的旋转框和真值旋转框做旋转 IoUIoU 大于 0.25 且角度差小于 15° 算一次正确预测。这样比单纯看 IoU 更严格也更接近真实机械臂抓取的容差范围。6.3 实际踩过的坑及解决方式第一个坑是角度回归的周期性问题。我最早直接在 FCN 里回归角度标量训练到第 5 个 epoch 时 loss 反复震荡角度输出永远差 90° 左右。后来改用“角度分类残差回归”结构周期性问题彻底消失。建议做旋转框相关任务时角度一律避开连续回归。第二个坑是旋转框 NMS 的工程实现。推理时多个旋转锚框会重叠必须做 NMS但旋转框的 IoU 计算比水平框慢很多直接导致推理速度降到 2 FPS。我的解决方式是用中心距离加角度差做近似 NMS如果两个旋转框中心距离小于一定阈值且角度差小于阈值就按质量分数保留一个。实测把推理速度拉回 10 FPS 以上准确率损失在可接受范围。第三个坑是 ROI 特征尺寸与 FCN 的感受野不匹配。一开始 ROI Align 输出 7×7 特征FCN 在这个小特征图上做预测抓取位姿的定位精度明显不足尤其在细长物体上。后来我把抓取分支的 ROI Align 输出调整为 14×14再在 FCN 内部通过堆叠小卷积逐步扩大感受野效果改善明显。第四个坑是关系类别不均衡。初期关系分类准确率只有 0.7后来加了类别权重和 Dropout直接到 0.84。如果读者在其它数据集上做关系推理建议第一时间看类别分布大部分关系数据集的类别分布都极度长尾。第五个坑是端到端训练不稳定。关系分支和抓取分支同时接在检测特征后面早期阶段检测框质量差两个分支的梯度同时回传容易互相干扰。我最后采用了两阶段训练策略先单独训练检测器 6 个 epoch然后冻结检测头单独训练抓取分支和关系分支 4 个 epoch最后全部解冻联合微调 2 个 epoch。整个过程稳定很多最终精度也更高。做这个项目最深的体会是多物体抓取真正的难点不在“抓”本身而在“理解和决策”。检测、抓取位姿、关系推理三个模块分开看都不算难难的是让它们共享特征、协同决策而不互相拖累。这套基于 PyTorch 的管线在 VMRD 上验证下来整体抓取成功率已经到了可以接机械臂做实物测试的水平。如果后续想扩展我会考虑在关系推理里引入时序信息让系统能根据抓取历史动态更新场景图而不是每帧独立推理。最后再分享一个小技巧调试时在关系可视化上多花点时间把关系图直接画在图像上比盯着一堆数值损失值直观得多很多问题一眼就能看出来。本文还有配套的精品资源点击获取

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价