本文主要讲的是一个名叫PlaneRecTR的基于transformer的模型,它能够从单幅图像当中恢复3D平面,完成平面相关的任务如平面检测,平面分割,平面参数评估,深度估计等。受到DETR的启发,模型也使用了可查询的query技术。
模型主要由三部分组成:
- 一个从输入的彩色图像当中学习密集的像素级深度嵌入的像素级模型
- 基于Transformer的统一查询学习模块,用于联合预测每个N个共享平面查询的四个目标属性,包括平面分类概率pi,平面参数ni,掩码嵌入和深度嵌入(i ∈ [1, 2, ..., N])。
- 平面级模块,通过掩码和深度嵌入(i ∈ [1, 2, ..., N]),生成平面级掩码mi和平面级深度di。
模型结构
模型的整体结构如上图所示。下面详细说一下上面的三个模块。
Pixel-Level Module
我们将输入的H X W的图像输入到预训练的ResNet-50当中来生成特征图。然后使用多尺度的卷积像素解码器来生成4个不同尺度的密集特征图,具体公式如下所示:
前三个特征图F1,F2,F3被输入到Transformer当中去。最后的特征图是一个
密集逐像素嵌入的特征图用来计算平面级的二进制掩码和平面级深度。
Transformer Module
这个模块使用的是带有mask attention的transformer解码器,从上面输入的三个不同尺度的特征图和N个可学习query来计算N个平面级别的嵌入,然后通过四个不同的多层感知器(MLP),将预测的嵌入独立地投影到四个目标属性上。
Plane-Level Module
我们通过第i个mask/depth embeding与Pixel-Level Module生成的密集做点积,生成密集的平面级二进制掩码和深度预测
。
3D Plane Recovery during Inference
在这个阶段会通过预测的数据来对平面进行重建。具体步骤是:先使用Pi来去除非平面的部分,得到具有K个平面的子集(K<=N)。然后通过计算
来找到最有可能的平面索引来作为最终的全局分割掩码。最后通过平面参数来推理平面的深度信息。
模型训练
Plane-level Depth Training.
与其他模型不同的是,我们想要让平面分割和几何推理实现互利。我们使用可学习的joint query以简单的多任务方式来统一平面恢复的所有部分。因此,我们明确地从共享特征空间中通过Transformer的注意机制来预测密集的平面级深度、二进制掩码、平面概率和平面参数。这种方法使得我们能够在一个统一的框架中同时处理平面的语义和几何信息,从而更好地推断和理解图像中的平面结构。
Bipartite Matching(二分匹配)
在训练的过程当中,我们需要让N个预测平面和M个真实平面形成最佳的匹配关系,所以需要用到二匹配。
Loss Function
损失函数的计算总和公式展示在下方:
λ是一个平衡权重,被设置为2,
其中
总结
本文介绍了以一个基于Transformer的模型,使用了可学习的query,可以处理像平面检测,平面分割,平面参数估计,平面深度估计等和平面相关的任务。