def projection_2d_sym(self, pose_pred, pose_targets, K, threshold=5):
model_2d_pred = pvnet_pose_utils.project(self.model, K, pose_pred)
model_2d_targets = pvnet_pose_utils.project(self.model, K,
pose_targets)
proj_mean_diff = np.mean(
find_nearest_point_distance(model_2d_pred, model_2d_targets))
self.proj2d.append(proj_mean_diff < threshold)
def projection_2d(self, pose_pred, pose_targets, K, threshold=5):
model_2d_pred = pvnet_pose_utils.project(self.model, K, pose_pred)
model_2d_targets = pvnet_pose_utils.project(self.model, K,
pose_targets)
proj_mean_diff = np.mean(
np.linalg.norm(model_2d_pred - model_2d_targets, axis=-1))
self.proj2d.append(proj_mean_diff < threshold)
def projection_2d(self, pose_pred, pose_targets, K, threshold=5):
if "obj_verts_2d" not in self.data.keys():
self.data["obj_verts_2d"] = []
model_2d_pred = pvnet_pose_utils.project(self.model, K, pose_pred)
model_2d_pred = np.expand_dims(model_2d_pred, axis=0) # add batch dimension
model_2d_targets = pvnet_pose_utils.project(self.model, K, pose_targets)
model_2d_targets = np.expand_dims(model_2d_targets, axis=0) # add batch dimension
proj_diff = np.linalg.norm(model_2d_pred - model_2d_targets, ord=2, axis=-1)
proj_diff = np.mean(proj_diff, axis=-1)
self.data["obj_verts_2d"].append(proj_diff)
self.proj2d.append(proj_diff < threshold)
def visualize(self, output, batch, id=0):
img = batch['img'][0].detach().cpu().numpy()
center = output['center'][0]
scale = output['scale'][0]
h, w = tless_pvnet_utils.input_scale
trans_output_inv = data_utils.get_affine_transform(center,
scale,
0, [w, h],
inv=1)
kpt = output['kpt_2d'].detach().cpu().numpy()
kpt_2d = data_utils.affine_transform(kpt, trans_output_inv)[0]
img_id = int(batch['img_id'][0])
anno = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))[0]
kpt_3d = np.concatenate([anno['fps_3d'], [anno['center_3d']]], axis=0)
K = np.array(anno['K'])
pose_gt = np.array(anno['pose'])
pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
corner_3d = np.array(anno['corner_3d'])
corner_2d_gt = pvnet_pose_utils.project(corner_3d, K, pose_gt)
corner_2d_pred = pvnet_pose_utils.project(corner_3d, K, pose_pred)
_, ax = plt.subplots(1)
ax.imshow(img)
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='b'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='b'))
plt.show()
def visualize(self, output, batch):
inp = img_utils.unnormalize_img(batch['inp'][0], mean,
std).permute(1, 2, 0)
kpt_2d = output['kpt_2d'][0].detach().cpu().numpy()
img_id = int(batch['img_id'][0])
anno = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))[0]
kpt_3d = np.concatenate([anno['fps_3d'], [anno['center_3d']]], axis=0)
K = np.array(anno['K'])
pose_gt = np.array(anno['pose'])
pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
corner_3d = np.array(anno['corner_3d'])
corner_2d_gt = pvnet_pose_utils.project(corner_3d, K, pose_gt)
corner_2d_pred = pvnet_pose_utils.project(corner_3d, K, pose_pred)
_, ax = plt.subplots(1)
ax.imshow(inp)
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='b'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='b'))
plt.show()
def obj_kpt_2d(self, pose_gt, kpt_3d, kpt_2d, K):
if "obj_keypoints_2d" not in self.data.keys():
self.data["obj_keypoints_2d"] = []
gt_kpt_2d = pvnet_pose_utils.project(kpt_3d, K, pose_gt)
#print("kpt_3d:\n{}".format(kpt_3d))
#print("K:\n{}".format(K))
#print("gt_kpt_2d:\n{}".format(gt_kpt_2d))
#print("kpt_2d:\n{}".format(kpt_2d))
err_2d = np.linalg.norm(gt_kpt_2d - kpt_2d, ord=2, axis=-1) # (K,1)
#print("err_2d:\n{}".format(err_2d))
err_2d = np.mean(err_2d, axis=-1)
#print("err_2d mean:\n{}".format(err_2d))
self.data["obj_keypoints_2d"].append(err_2d)
def visualize_demo(self, output, inp, meta):
inp = img_utils.unnormalize_img(inp[0], mean, std).permute(1, 2, 0)
kpt_2d = output['kpt_2d'][0].detach().cpu().numpy()
kpt_3d = np.array(meta['kpt_3d'])
K = np.array(meta['K'])
pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
corner_3d = np.array(meta['corner_3d'])
corner_2d_pred = pvnet_pose_utils.project(corner_3d, K, pose_pred)
_, ax = plt.subplots(1)
ax.imshow(inp)
ax.add_patch(patches.Polygon(xy=corner_2d_pred[[0, 1, 3, 2, 0, 4, 6, 2]], fill=False, linewidth=1, edgecolor='b'))
ax.add_patch(patches.Polygon(xy=corner_2d_pred[[5, 4, 6, 7, 5, 1, 3, 7]], fill=False, linewidth=1, edgecolor='b'))
plt.show()
def visualize(self, output, batch, name, high_resolution=None):
inp = img_utils.unnormalize_img(batch['inp'][0], mean,
std).permute(1, 2, 0)
#inpnew = high_resolution
kpt_2d = output['kpt_2d'][0].detach().cpu().numpy()
#print(kpt_2d)
img_id = int(batch['img_id'][0])
anno = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))[0]
kpt_3d = np.concatenate([anno['fps_3d'], [anno['center_3d']]], axis=0)
#K = np.array(anno['K'])
#print(K)
#print(K.shape)
K = np.array([[1.2430691e+03, 0, 5.45181403e+02],
[0, 1.23942895e+03, 7.0817400e+02], [0, 0, 1]])
#pose_gt = np.array(anno['pose'])
pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
corner_3d = np.array(anno['corner_3d'])
#corner_2d_gt = pvnet_pose_utils.project(corner_3d, K, pose_gt)
corner_2d_pred = pvnet_pose_utils.project(corner_3d, K, pose_pred)
#print(corner_2d_pred)
#corner_2d_pred = corner_2d_pred * 2.44
_, ax = plt.subplots(1)
plt.axis('off')
frame = plt.gca()
frame.axes.get_yaxis().set_visible(False)
frame.axes.get_xaxis().set_visible(False)
ax.imshow(inp)
#ax.add_patch(patches.Polygon(xy=corner_2d_gt[[0, 1, 3, 2, 0, 4, 6, 2]], fill=False, linewidth=1, edgecolor='g'))
#ax.add_patch(patches.Polygon(xy=corner_2d_gt[[5, 4, 6, 7, 5, 1, 3, 7]], fill=False, linewidth=1, edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=2,
edgecolor='r'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=2,
edgecolor='r'))
ax.figure.savefig(name, bbox_inches='tight', dpi=244)
def visualize(self, output, batch, output_file=None):
col_nb = 4
fig = plt.figure(figsize=(12, 3))
axes = fig.subplots(1, 4, gridspec_kw={"left": 0.0, "right": 1.0, "bottom": 0.0, "top": 1.0})
inp = img_utils.unnormalize_img(batch['inp'][0], mean, std).permute(1, 2, 0)
kpt_2d = output['kpt_2d'][0].detach().cpu().numpy()
var_2d = output['var'][0].detach().cpu().numpy()
img_id = int(batch['img_id'][0])
anno = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))[0]
kpt_3d = np.concatenate([anno['fps_3d'], [anno['center_3d']]], axis=0)
K = np.array(anno['K'])
pose_gt = np.array(anno['pose'])
#pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
pose_pred = pvnet_pose_utils.uncertainty_pnp(kpt_3d, kpt_2d, var_2d, K)
kpt_2d_gt = pvnet_pose_utils.project(kpt_3d, K, pose_gt)
verts_2d_gt = pvnet_pose_utils.project(np.array(self.mesh.vertices, copy=True), K, pose_gt)
verts_2d_pred = pvnet_pose_utils.project(np.array(self.mesh.vertices, copy=True), K, pose_pred)
rot_verts_3d_gt = np.dot(np.array(self.mesh.vertices, copy=True), pose_gt[:, :3].T) + pose_gt[:, 3:].T
rot_verts_3d_pred = np.dot(np.array(self.mesh.vertices, copy=True), pose_pred[:, :3].T) + pose_pred[:, 3:].T
# Column 0
col_idx = 0
axes[col_idx].imshow(inp)
axes[col_idx].axis("off")
if kpt_2d_gt is not None:
axes[col_idx].scatter(
kpt_2d_gt[:, 0], kpt_2d_gt[:, 1], c="b", s=2, marker="X", alpha=0.7
)
if kpt_2d is not None:
axes[col_idx].scatter(
kpt_2d[:, 0], kpt_2d[:, 1], c="r", s=2, marker="X", alpha=0.7
)
if var_2d is not None:
ells = compute_confidence_ellipses(kpt_2d, var_2d)
for e in ells:
axes[col_idx].add_artist(e)
if kpt_2d_gt is not None and kpt_2d is not None:
arrow_nb = kpt_2d_gt.shape[0]
idxs = range(arrow_nb)
arrows = np.concatenate([kpt_2d_gt[idxs].astype(np.float), kpt_2d[idxs].astype(np.float)], axis=0)
links = [[i, i + arrow_nb] for i in idxs]
_visualize_joints_2d(
axes[col_idx],
arrows,
alpha=0.5,
joint_idxs=False,
links=links,
color=["k"] * arrow_nb,
)
# Column 1
col_idx = 1
axes[col_idx].imshow(inp)
axes[col_idx].axis("off")
# Visualize 2D object vertices
if verts_2d_gt is not None:
axes[col_idx].scatter(
verts_2d_gt[:, 0], verts_2d_gt[:, 1], c="b", s=1, alpha=0.05
)
if verts_2d_pred is not None:
axes[col_idx].scatter(
verts_2d_pred[:, 0], verts_2d_pred[:, 1], c="r", s=1, alpha=0.02
)
# Column 2
# view from the top
col_idx = 2
if rot_verts_3d_gt is not None:
axes[col_idx].scatter(
rot_verts_3d_gt[:, 2], rot_verts_3d_gt[:, 0], c="b", s=1, alpha=0.02
)
if rot_verts_3d_pred is not None:
axes[col_idx].scatter(
rot_verts_3d_pred[:, 2], rot_verts_3d_pred[:, 0], c="r", s=1, alpha=0.02
)
axes[col_idx].invert_yaxis()
# Column 3
# view from the right
col_idx = 3
# invert second axis here for more consistent viewpoints
if rot_verts_3d_gt is not None:
axes[col_idx].scatter(
rot_verts_3d_gt[:, 2], -rot_verts_3d_gt[:, 1], c="b", s=1, alpha=0.02
)
if rot_verts_3d_pred is not None:
axes[col_idx].scatter(
rot_verts_3d_pred[:, 2], -rot_verts_3d_pred[:, 1], c="r", s=1, alpha=0.02
)
_squashfig(fig)
if output_file is None:
plt.show()
else:
fig.savefig(output_file, dpi=300)
plt.close()
def visualize(self, output, batch):
img_id = int(batch['meta']['img_id'])
img_data = self.coco.loadImgs(int(img_id))[0]
path = img_data['file_name']
depth_path = img_data['depth_path']
img = np.array(Image.open(path))
ann_ids = self.coco.getAnnIds(imgIds=img_id, catIds=self.obj_id)
annos = self.coco.loadAnns(ann_ids)
kpt_3d = np.concatenate([annos[0]['fps_3d'], [annos[0]['center_3d']]],
axis=0)
corner_3d = np.array(annos[0]['corner_3d'])
K = np.array(annos[0]['K'])
kpt_2d = output['kpt_2d'].detach().cpu().numpy()
centers = batch['meta']['center']
scales = batch['meta']['scale']
boxes = batch['meta']['box']
h, w = batch['inp'].size(2), batch['inp'].size(3)
kpt_2ds = []
segs = []
for i in range(len(centers)):
center = centers[i].detach().cpu().numpy()
scale = scales[i].detach().cpu().numpy()
kpt_2d_ = kpt_2d[i]
trans_inv = data_utils.get_affine_transform(center[0],
scale[0],
0, [w, h],
inv=1)
kpt_2d_ = data_utils.affine_transform(kpt_2d_, trans_inv)
kpt_2ds.append(kpt_2d_)
seg = torch.argmax(output['seg'][i], dim=0).detach().cpu().numpy()
seg = seg.astype(np.uint8)
seg = cv2.warpAffine(seg,
trans_inv, (720, 540),
flags=cv2.INTER_NEAREST)
segs.append(seg)
_, ax = plt.subplots(1)
ax.imshow(img)
# for i in range(len(boxes)):
# x_min, y_min, x_max, y_max = boxes[i].view(-1).numpy()
# ax.plot([x_min, x_min, x_max, x_max, x_min], [y_min, y_max, y_max, y_min, y_min])
depth = np.array(Image.open(depth_path)).astype(np.float32)
for i, kpt_2d in enumerate(kpt_2ds):
pose_pred = pvnet_pose_utils.pnp(kpt_3d, kpt_2d, K)
mask = segs[i]
box = cv2.boundingRect(mask.astype(np.uint8))
x, y = box[0] + box[2] / 2., box[1] + box[3] / 2.
z = np.mean(depth[mask != 0] / 10000.)
x = ((x - K[0, 2]) * z) / float(K[0, 0])
y = ((y - K[1, 2]) * z) / float(K[1, 1])
center = [x, y, z]
# pose_pred[:, 3] = center
corner_2d_pred = pvnet_pose_utils.project(corner_3d, K, pose_pred)
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='b'))
ax.add_patch(
patches.Polygon(xy=corner_2d_pred[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='b'))
for anno in annos:
pose_gt = np.array(anno['pose'])
corner_2d_gt = pvnet_pose_utils.project(corner_3d, K, pose_gt)
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[0, 1, 3, 2, 0, 4, 6, 2]],
fill=False,
linewidth=1,
edgecolor='g'))
ax.add_patch(
patches.Polygon(xy=corner_2d_gt[[5, 4, 6, 7, 5, 1, 3, 7]],
fill=False,
linewidth=1,
edgecolor='g'))
plt.show()
