mean = pvnet_config.mean
std = pvnet_config.std
class Visualizer:
def __init__(self):
args = DatasetCatalog.get(cfg.test.dataset)
self.ann_file = args['ann_file']
self.coco = coco.COCO(self.ann_file)
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()我想在画出边界框的同时画出关键点kpt_2d的位置
最新发布