标题:PoseCNN: A Convolutional Neural Network for 6D Object Pose Estimation in Cluttered Scenes
作者团队:NVIDIA, 华盛顿大学
期刊会议
时间:2018
代码https://rse-lab.cs.washington.edu/projects/posecnn/

一、介绍

1.1 问题提出

物体之间的遮挡和聚集增加了6D位姿估计的难度。

现有方法

  • 传统基于RGB图像,依赖纹理,无法处理弱纹理物体
  • 基于RGBD模板,受遮挡影响较大
  • 2D->3D的方法无法处理对称物体

1.2 方法介绍

(1)概述

PoseCNN(利用RGB图像进行位姿估计,使用RGBD图像进行位姿优化)

  • 物体的3D位置和到相机的距离通过计算图像中物体的中心来估计。
  • 物体的3D旋转通过回归到四元数来估计。

提出了一种损失函数解决对称物体的估计。

(2)思路

PoseCNN的关键想法:将位姿估计解耦为位置和姿态两个不同的部分,使网络可以对两个系统的独立性和依赖性建模。

  • 预测每个像素的语义标签
  • 估计2D图像坐标系下的物体中心
  • 估计物体中心与相机的距离(利用每个像素语义标签和物体中心位置)
  • 估计3D位置
  • 估计3D姿态(通过边界框内物体的四元数)

(3)对称物体的处理

由于对称物体在不同角度观测时可能得到同样的结果,无法估计它们的姿态。

提出ShapeMatch-Loss。

二、PoseCNN

物体的位姿估计,就是估计物体坐标系与相机坐标系之间的刚性变换SE3,包含旋转R和平移T。T决定了图像中物体的位置和大小,R决定了物体在图像中的形状和纹理。

2.1 网络简述

PoseCNN主要分为两个阶段:

  1. 主干网络:由13个卷积层和4个最大池化层组成,提取特征图。
  2. 将主干网络生成的高为特征图映射到低维。
  3. 进行三个不同的任务估计

2.2 语义标签

语义标签分支使用了两张特征图,通道数均为512,分辨率分别从原始图像降低了1/8和1/16。

该分支首先将特征图维度降为64,然后使用反卷积层将1/16图像的分辨率翻倍,获得两个64通道,1/8分辨率的特征图。

然后将两个特征图相加,在进行反卷积,将分辨率提高8倍,获得与原始图象大小相同的特征图。

最后对特征图进行卷积层运算,生成像素的语义标签分数。该层的输出有n个通道(n为语义分类数)。训练时使用Softmax交叉商损失函数训练语义标签分支,测试是使用softmax函数计算像素的类概率。

2.3 3D位置估计

(1)原理

通过定位物体在图像中的中心和与相机的距离,估计3D平移。

原理是,只要获得T在图像上的投影c=(cx,cy)Tc=(c_x, c_y)^T,以及距离TzT_z,就可以计算出TxT_xTyT_y

其中fxf_xfyf_y是焦距,pxp_xpyp_y是像主点。

(2)确定物体中心

常规的方法是直接检测中心点,但这样的问题是无法处理物体被遮挡的情况。

本文利用隐式形状模型的思想,将图像上的每一个像素,回归为3个变量:

其中P=(x,y)P=(x, y)是图像中的一个像素点,(cx,cy)(c_x, c_y)是物体中心座标点。nxn_x表示xx方向上,像素点与中心点之间的偏移,并进行归一化处理。nyn_y表示y方向上的偏移,TzT_z表示像素点的预测深度。

因此需要预测的变量包括(nx,ny,Tz)(n_x, n_y, T_z)

由于需要回归三个变量,因此使用了128通道的特征图,最后一个卷积层有3xn通道,n为物体类别数。

得到了每个像素预测的3个变量后,需要确定物体中心坐标

引入了霍夫投票机制,将像素级 语义标签结果 和 中心回归结果 作为输入。

每个像素沿网络预测的射线方向对图像中心位置进行投票。当所有物体类中的像素完成投票后,可以获得所有的投票分数,选取分数最高的位置即为物体中心位置。

对于同一图像中有多个同类物体的情况,使用非极大抑制,通过设定阈值来实现分数选择。

通过以上操作,就可以确定物体中心的三维坐标,其中预测的TzT_z认为是物体内部中心点的距离。

2.4 3D姿态估计

(1)网络结构

3D旋转估计分支利用了主干网络得到的特征图的ROI区域部分,通过两个池化层,来实现3D回归。

将所有特征图加在一起,送入三个全连接层中,前两个全连接层维度为4096,最后一个全连接层维度为4xn,其中n为对象类别数,每个对象输出以四元数表示的3D旋转。

(2)损失函数的设计

为了回归四元数,设计了两个损失函数。

PoseLoss(PLoss)

在3维空间中计算,计算正确模型位姿的点与测量的位姿的点的平均平方距离。

其中MM代表一系列三维模型的点,m为点的个数,R(q)R(\overline q)R(q)R(q)分别代表计算的旋转矩阵和真实的旋转矩阵。

ShapeMatch-Loss(SLOSS)

由于PLoss损失函数无法处理对称物体的情况,对于对陈物体来说,会有很多个正确的三维旋转矩阵。

因此设计了损失函数如下:

该损失函数测量的是预测模型上的点与真实模型的最近的点的偏移量。通过这种方式,SLoss不会惩罚对称物体绕对称轴的等效旋转。

注意:是计算法中不需要在三维空间中计算距离,而是将变换后的点投影到图像上,计算图像空间中的距离(重投影误差)。

三、实验

3.1 数据集

在YCB数据集和Occluded-LINEMOD数据集(在LINEMOD数据集中挑选一段物体严重遮挡的视频)进行测试。

3.2 对比基准

由于文章发布时主流的位姿估计方法是将图像像素回归到3D物体坐标。因此作者实现了一个根据PoseCNN网络修改的变体,在网络中直接回归每个像素在物体坐标系的三维坐标,移除了3D旋转预测分支,而是利用语义标签结果和三维物体坐标结果,通过RANSAC恢复6D位姿。

3.3 改进

当深度可用时,使用ICP算法来优化6D姿态。

四、PoseCNN可视化

以下内容原作者为Panpanpan!
分类专栏: 【DenseFusion进阶】 文章标签: 深度学习 pytorch 计算机视觉
版权声明:本文为博主原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。
本文链接:https://blog.csdn.net/weixin_44564705/article/details/125129844
DenseFusion系列代码全讲解目录:【DenseFusion系列目录】代码全讲解+可视化+计算评估指标_Panpanpan!的博客-CSDN博客
这些内容均为个人学习记录,欢迎大家提出错误一起讨论一起学习!


可视化的效果展示

选择要可视化的图片,输入路径,然后对其进行eval,估计姿态pred,输出target点云和pred点云在二维图像上的投影,如下:

左边的图为pred,右边的图为target。

具体操作

代码放在了tools/draw_linemod.py

脚本放在了experiments/scripts/draw_linemod.sh

先看脚本文件内容:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#!/bin/bash

set -x
set -e

export PYTHONUNBUFFERED="True"
export CUDA_VISIBLE_DEVICES=0

python3 ./tools/draw_linemod.py --data_root ./datasets/linemod/Linemod_preprocessed/data/08\
--item 0707\
--seg_root ./datasets/linemod/Linemod_preprocessed/segnet_results/08_label\
--obj 8\
--model_root ./datasets/linemod/Linemod_preprocessed/models\
--model trained_models/linemod/original/pose_model_current.pth\
--refine_model trained_models/linemod/original/pose_refine_model_current.pth\
--output ./

参数分别为:

–data_root:要选择的图像类别文件夹路径
  --item:选择的图像编号
  --seg_root:因为是eval模式,需要用到标准分割的标签,这里输入选择的图像类别的语义分割路径
  --obj:图像类别编号
  --model_root:图像类别元数据模型路径
  --model:训练好的PoseNet路径
  --refine_model:训练好的PoseRefineNet路径

–output:可视化保存路径

输入这些参数,然后在DenseFusion文件夹下执行:

1
./experiments/scripts/draw_linemod.sh

如果是最开始执行,可能会提示.sh文件没有权限,就要执行以下命令打开权限:

1
chmod +x ./experiments/scripts/draw_linemod.sh

然后就能在保存路径里查看这两个图片。

完整代码

draw_linemod.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
# -*- coding: utf-8 -*-
import _init_paths
import argparse
import os
from PIL import Image
import numpy as np
import numpy.ma as ma
import torch
import random
import torchvision.transforms as transforms
import yaml
import cv2
from torch.autograd import Variable
from lib.network import PoseNet, PoseRefineNet
from lib.loss import Loss
from lib.loss_refiner import Loss_refine
from lib.knn.__init__ import KNearestNeighbor
from lib.transformations import euler_matrix, quaternion_matrix, quaternion_from_matrix
import copy
import os
from PIL import Image
import matplotlib.pyplot as plt

border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
img_width = 480
img_length = 640
def mask_to_bbox(mask):
mask = mask.astype(np.uint8)
_, contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

x = 0
y = 0
w = 0
h = 0
for contour in contours:
tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contour)
if tmp_w * tmp_h > w * h:
x = tmp_x
y = tmp_y
w = tmp_w
h = tmp_h
return [x, y, w, h]


def get_bbox(bbox):
bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
if bbx[0] < 0:
bbx[0] = 0
if bbx[1] >= 480:
bbx[1] = 479
if bbx[2] < 0:
bbx[2] = 0
if bbx[3] >= 640:
bbx[3] = 639
rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]
r_b = rmax - rmin
for tt in range(len(border_list)):
if r_b > border_list[tt] and r_b < border_list[tt + 1]:
r_b = border_list[tt + 1]
break
c_b = cmax - cmin
for tt in range(len(border_list)):
if c_b > border_list[tt] and c_b < border_list[tt + 1]:
c_b = border_list[tt + 1]
break
center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
rmin = center[0] - int(r_b / 2)
rmax = center[0] + int(r_b / 2)
cmin = center[1] - int(c_b / 2)
cmax = center[1] + int(c_b / 2)
if rmin < 0:
delt = -rmin
rmin = 0
rmax += delt
if cmin < 0:
delt = -cmin
cmin = 0
cmax += delt
if rmax > 480:
delt = rmax - 480
rmax = 480
rmin -= delt
if cmax > 640:
delt = cmax - 640
cmax = 640
cmin -= delt
return rmin, rmax, cmin, cmax
def ply_vtx(path):
f = open(path)
assert f.readline().strip() == "ply"
f.readline()
f.readline()
N = int(f.readline().split()[-1])
while f.readline().strip() != "end_header":
continue
pts = []
for _ in range(N):
pts.append(np.float32(f.readline().split()[:3]))
return np.array(pts)

def get_data(data_root, item, seg_root, obj, model_root, add_noise, noise_trans):
img = Image.open('{0}/rgb/{1}.png'.format(data_root, item))
ori_img = np.array(img)
depth = np.array(Image.open('{0}/depth/{1}.png'.format(data_root, item)))
label = np.array(Image.open('{0}/{1}_label.png'.format(seg_root, item)))
print('{0}/depth/{1}.png'.format(data_root, item))
print('{0}/{1}_label.png'.format(seg_root, item))
print('{0}/gt.yml'.format(data_root))
print('{0}/obj_{1}.ply'.format(model_root, '%02d' % obj))

meta_file = open('{0}/gt.yml'.format(data_root), 'r')
meta = yaml.load(meta_file, Loader=yaml.FullLoader)
pt = ply_vtx('{0}/obj_{1}.ply'.format(model_root, '%02d' % obj))

trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
num = 500
xmap = np.array([[j for i in range(640)] for j in range(480)])
ymap = np.array([[i for i in range(640)] for j in range(480)])
cam_cx = 325.26110
cam_cy = 242.04899
cam_fx = 572.41140
cam_fy = 573.57043
num_pt_mesh_large = 500
num_pt_mesh_small = 500
rank = int(item)
if obj == 2:
for i in range(0, len(meta[obj])):
if meta[rank][i]['obj_id'] == 2:
meta = meta[rank][i]
break
else:
meta = meta[rank][0]

mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
mask_label = ma.getmaskarray(ma.masked_equal(label, np.array(255)))
mask = mask_label * mask_depth

if add_noise:
img = trancolor(img)

img = np.array(img)[:, :, :3]
img = np.transpose(img, (2, 0, 1))
img_masked = img

rmin, rmax, cmin, cmax = get_bbox(mask_to_bbox(mask_label))
img_masked = img_masked[:, rmin:rmax, cmin:cmax]
#p_img = np.transpose(img_masked, (1, 2, 0))
#scipy.misc.imsave('evaluation_result/{0}_input.png'.format(index), p_img)

target_r = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
target_t = np.array(meta['cam_t_m2c'])
add_t = np.array([random.uniform(-noise_trans, noise_trans) for i in range(3)])

choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
if len(choose) == 0:
cc = torch.LongTensor([0])
return(cc, cc, cc, cc, cc, cc)

if len(choose) > num:
c_mask = np.zeros(len(choose), dtype=int)
c_mask[:num] = 1
np.random.shuffle(c_mask)
choose = choose[c_mask.nonzero()]
else:
choose = np.pad(choose, (0, num - len(choose)), 'wrap')

depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
xmap_masked = xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
ymap_masked = ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
choose = np.array([choose])

cam_scale = 1.0
pt2 = depth_masked / cam_scale
pt0 = (ymap_masked - cam_cx) * pt2 / cam_fx
pt1 = (xmap_masked - cam_cy) * pt2 / cam_fy
cloud = np.concatenate((pt0, pt1, pt2), axis=1)
cloud = cloud / 1000.0
print(cloud.shape)

if add_noise:
cloud = np.add(cloud, add_t)

#fw = open('evaluation_result/{0}_cld.xyz'.format(index), 'w')
#for it in cloud:
# fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
#fw.close()
model_points = pt/ 1000.0
dellist = [j for j in range(0, len(model_points))]
dellist = random.sample(dellist, len(model_points) - num_pt_mesh_small)
model_points = np.delete(model_points, dellist, axis=0)

#fw = open('evaluation_result/{0}_model_points.xyz'.format(index), 'w')
#for it in model_points:
# fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
#fw.close()

target = np.dot(model_points, target_r.T)
if add_noise:
target = np.add(target, target_t / 1000.0 + add_t)
out_t = target_t / 1000.0 + add_t
else:
target = np.add(target, target_t / 1000.0)
out_t = target_t / 1000.0

#fw = open('evaluation_result/{0}_tar.xyz'.format(index), 'w')
#for it in target:
# fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
#fw.close()

return torch.from_numpy(cloud.astype(np.float32)), \
torch.LongTensor(choose.astype(np.int32)), \
norm(torch.from_numpy(img_masked.astype(np.float32))), \
torch.from_numpy(target.astype(np.float32)), \
torch.from_numpy(model_points.astype(np.float32)), \
torch.LongTensor([obj])


parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default = '', help='ycb or linemod')
parser.add_argument('--seg_root', type=str, default = '', help='ycb or linemod')
parser.add_argument('--model_root', type=str, default = '', help='ycb or linemod')
parser.add_argument('--item', type=str, default = '', help='')
parser.add_argument('--obj', type=int, default=1, help='')
parser.add_argument('--model', type=str, default = '', help='resume PoseNet model')
parser.add_argument('--refine_model', type=str, default = '', help='resume PoseRefineNet model')
parser.add_argument('--output', type=str, default = '', help='resume PoseRefineNet model')
opt = parser.parse_args()


estimator = PoseNet(num_points = 500, num_obj = 13)
estimator.cuda()
refiner = PoseRefineNet(num_points = 500, num_obj = 13)
refiner.cuda()
estimator.load_state_dict(torch.load(opt.model))
refiner.load_state_dict(torch.load(opt.refine_model))

bs = 1
num_points = 500
iteration = 4
knn = KNearestNeighbor(1)

for i in range(1):
points, choose, img, target, model_points, idx = get_data(opt.data_root, opt.item, opt.seg_root, opt.obj, opt.model_root, True, 0.0)

if len(points.size()) == 1:
print('No.{0} NOT Pass! Lost detection!'.format(i))
continue
points, choose, img, target, model_points, idx = Variable(points.unsqueeze(0)).cuda(), \
Variable(choose.unsqueeze(0)).cuda(), \
Variable(img.unsqueeze(0)).cuda(), \
Variable(target.unsqueeze(0)).cuda(), \
Variable(model_points.unsqueeze(0)).cuda(), \
Variable(idx.unsqueeze(0)).cuda()
pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, num_points, 1)
pred_c = pred_c.view(bs, num_points)
how_max, which_max = torch.max(pred_c, 1)
pred_t = pred_t.view(bs * num_points, 1, 3)

my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
my_t = (points.view(bs * num_points, 1, 3) + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
my_pred = np.append(my_r, my_t)
for ite in range(0, iteration):
Tt = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(num_points, 1).contiguous().view(1, num_points, 3)
my_mat = quaternion_matrix(my_r)
R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3)
my_mat[0:3, 3] = my_t

new_points = torch.bmm((points - Tt), R).contiguous()
pred_r, pred_t = refiner(new_points, emb, idx)
pred_r = pred_r.view(1, 1, -1)
pred_r = pred_r / (torch.norm(pred_r, dim=2).view(1, 1, 1))
my_r_2 = pred_r.view(-1).cpu().data.numpy()
my_t_2 = pred_t.view(-1).cpu().data.numpy()
my_mat_2 = quaternion_matrix(my_r_2)
my_mat_2[0:3, 3] = my_t_2

my_mat_final = np.dot(my_mat, my_mat_2)
my_r_final = copy.deepcopy(my_mat_final)
my_r_final[0:3, 3] = 0
my_r_final = quaternion_from_matrix(my_r_final, True)
my_t_final = np.array([my_mat_final[0][3], my_mat_final[1][3], my_mat_final[2][3]])

my_pred = np.append(my_r_final, my_t_final)
my_r = my_r_final
my_t = my_t_final
model_points = model_points[0].cpu().detach().numpy()
my_r = quaternion_matrix(my_r)[:3, :3]
pred = np.dot(model_points, my_r.T) + my_t
target = target[0].cpu().detach().numpy()

pred = pred*1000.0
target = target*1000.0
cam_scale = 1.0
cam_cx = 325.26110
cam_cy = 242.04899
cam_fx = 572.41140
cam_fy = 573.57043

depth_masked = pred[:,2] * cam_scale
ymap_masked_pred = pred[:,0] * cam_fx / pred[:,2] + cam_cx
xmap_masked_pred = pred[:,1] * cam_fy / pred[:,2] + cam_cy

depth_masked = target[:,2] * cam_scale
ymap_masked_target = target[:,0] * cam_fx / target[:,2] + cam_cx
xmap_masked_target = target[:,1] * cam_fy / target[:,2] + cam_cy

image = Image.open('{0}/rgb/{1}.png'.format(opt.data_root, opt.item))

plt.figure(figsize=(8,12)) # 图像窗口名称
plt.imshow(image)
plt.scatter(ymap_masked_pred,xmap_masked_pred,marker='.',c='r',alpha=1)
plt.axis('on') # 关掉坐标轴为 off
plt.title('image') # 图像题目
plt.show()
plt.savefig('{0}/test_pred.png'.format(opt.output))

plt.figure(figsize=(8,12)) # 图像窗口名称
plt.imshow(image)
plt.scatter(ymap_masked_target,xmap_masked_target,marker='.',c='r',alpha=1)
plt.axis('on') # 关掉坐标轴为 off
plt.title('image') # 图像题目
plt.show()
plt.savefig('{0}/test_target.png'.format(opt.output))