def draw_corr_torch(img0, pix_pos0, img1, pix_pos1):
"""
Draw matches, torch version. Parameters are the same as draw keypoints
"""
img0, img1 = [tonumpy(x) for x in [img0, img1]]
img0, img1 = [touint8(x).astype(np.uint8) for x in [img0, img1]]
pix_pos0, pix_pos1 = [x.numpy() for x in [pix_pos0, pix_pos1]]
img = draw_corr(img0, pix_pos0, img1, pix_pos1)
img = totensor(tofloat(img))
return img
def draw_kps_torch(img0, pix_pos0, img1, pix_pos1):
"""
Draw keypoints, torch version
:param img0: (C, H, W)
:param pix_pos0: (N, 2), (x, y)
:param img1: (C, H, W)
:param pix_pos1: (N, 2), (x, y)
:return: a single torch image
"""
img0, img1 = [tonumpy(x) for x in [img0, img1]]
img0, img1 = [touint8(x) for x in [img0, img1]]
pix_pos0, pix_pos1 = [x.numpy() for x in [pix_pos0, pix_pos1]]
img = draw_kps(img0, pix_pos0, img1, pix_pos1)
img = totensor(tofloat(img))
return img
def draw_scale_torch(scale):
"""
scale: [H, W]
"""
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, figsize=(2.4, 2.4))
mappable = ax.imshow(scale.detach().cpu().numpy())
fig.colorbar(mappable, ax=ax)
fig.canvas.draw()
w, h = fig.canvas.get_width_height()
buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8).reshape(h, w, 4)
scale = np.roll(buf, 3, axis=2)[:, :, :3]
plt.close()
return totensor(tofloat(scale))
def draw_paired_desc_torch(desc0, kps0, image0, desc1, kps1, image1, H):
"""
desc0: [C, H_0', W_0']
kps0: [N, 2]
image0: [3, H_0, W_0]
desc1: [C, H_1', W_1']
kps1: [N, 2]
image1: [3, H_1, W_1]
H: [3, 3]
"""
desc0 = (desc2RGB(tonumpy(desc0)) * 255.).astype(np.uint8) # [H, W, 3]
kps0 = kps0.detach().cpu().numpy()
desc1 = (desc2RGB(tonumpy(desc1)) * 255.).astype(np.uint8)
kps1 = kps1.detach().cpu().numpy()
# compute the corresponding region
H = H.detach().cpu().numpy()
h1, w1 = image1.shape[1:]
pts1 = np.array(
[[0, 0],
[0, h1],
[w1, h1],
[w1, 0]]
).astype(np.float32)
pts0 = cv2.perspectiveTransform(np.reshape(pts1, [1, -1, 2]), np.linalg.inv(H))[0]
# compute the ratio between desc and image, and rescale the kps
ratio_h0 = desc0.shape[0] / image0.shape[1] # [H_0', W_0', 3], [3, H_0, W_0]
ratio_w0 = desc0.shape[1] / image0.shape[2]
pts0 *= np.array([[ratio_w0, ratio_h0]])
ratio_h1 = desc1.shape[0] / image1.shape[1]
ratio_w1 = desc1.shape[1] / image1.shape[2]
pts1 *= np.array([[ratio_w1, ratio_h1]])
# draw the corresponding region on the image
pts0 = pts0.astype(np.int32)
desc0 = cv2.polylines(desc0.copy(), [pts0.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
pts1 = pts1.astype(np.int32)
desc1 = cv2.polylines(desc1.copy(), [pts1.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
# draw correspondences
desc = draw_corr(desc0, kps0, desc1, kps1, num=10)
desc = totensor(tofloat(desc))
return desc
def get_tensorboard_data(self, num_kps=20):
"""
This processes the data needed for visualization. It expects the follow-
ing in self.logger
- image0: (C, H, W), Tensor, normalized
- image1: (C, H, W), Tensor, normalized
- desc0: (C, H, W)
- desc1: (C, H, W)
- kps0: (N, 2), each being (x, y)
- kps1: (N, 2)
- kps2: (N, 2), negative ones
And it returns a dictionary
- desc0: descriptor 1, RGB, (3, H, W)
- desc1: descriptor 2, RGB, (3, H, W)
- img0: image 1, (3, H, W)
- img1: image 2, (3, H, W)
- keypoints: the two images marked with num_kps keypoints
- neg_keypoints: image 2 marked with negative keypoints
- corr: ground truth correspondences
- corr false: false correspondences
"""
# original images
image0 = self.logger['image0']
image1 = self.logger['image1']
# descriptors
desc0 = self.logger['desc0']
desc1 = self.logger['desc1']
# keypoints
kps0 = self.logger['kps0']
kps1 = self.logger['kps1']
kps2 = self.logger['kps2']
# process the images
image0 = unnormalize(image0)
image1 = unnormalize(image1)
# process the descriptor
desc0, desc1 = [desc2RGB(tonumpy(x)) for x in [desc0, desc1]]
desc0, desc1 = [totensor(d) for d in [desc0, desc1]]
# choose keypoints
N = kps0.shape[0]
indices = np.random.choice(N, size=num_kps, replace=False)
# draw keypoints
kps = draw_kps_torch(image0, kps0[indices], image1, kps1[indices])
# draw negative keypoints
neg_kps = draw_kps_torch(image0, kps0[indices], image1, kps2[indices])
# draw correspondences
corr_gt = draw_corr_torch(image0, kps0[indices], image1, kps1[indices])
# draw correspondences
corr_false = draw_corr_torch(image0, kps0[indices], image1,
kps2[indices])
return {
'img0': image0,
'img1': image1,
'desc0': desc0,
'desc1': desc1,
'keypoints': kps,
'neg_keypoints': neg_kps,
'corr': corr_gt,
'corr_false': corr_false
}