def draw_corspd_region_torch(img0, img1, H):
"""
img0: [3, H_0, W_0]
img1: [3, H_1, W_1]
"""
img0 = (tonumpy(unnormalize(img0)) * 255.).astype(np.uint8) # [H0, W0, 3]
img1 = (tonumpy(unnormalize(img1)) * 255.).astype(np.uint8) # [H1, W1, 3]
if isinstance(H, torch.Tensor):
H = H.detach().cpu().numpy()
h1, w1 = img1.shape[:2]
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]
# draw the corresponding region on the image
pts0 = pts0.astype(np.int32)
img0 = cv2.polylines(img0.copy(), [pts0.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
pts1 = pts1.astype(np.int32)
img1 = cv2.polylines(img1.copy(), [pts1.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
return np.concatenate([img0, img1], axis=1)
def draw_img_desc_torch(img, desc):
"""
img: [3, H, W]
desc: [C, H, W]
"""
img = (tonumpy(unnormalize(img)) * 255.).astype(np.uint8) # [H, W, 3]
desc = (desc2RGB(tonumpy(desc)) * 255.).astype(np.uint8) # [H, W, 3]
h, w = img.shape[:2]
desc = cv2.resize(desc.copy(), (w, h), interpolation=cv2.INTER_LINEAR)
return np.concatenate([img, desc], axis=1)
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 draw_paired_img_desc_torch(img0, desc0, kps0, img1, desc1, kps1, H):
"""
img0: [3, H_0, W_0]
desc0: [C, H_0', W_0']
kps0: [N, 2]
img1: [3, H_1, W_1]
desc1: [C, H_1', W_1']
kps1: [N, 2]
H: [3, 3]
"""
img0 = (tonumpy(unnormalize(img0)) * 255.).astype(np.uint8) # [H, W, 3]
desc0 = (desc2RGB(tonumpy(desc0)) * 255.).astype(np.uint8) # [H, W, 3]
kps0 = kps0.detach().cpu().numpy()
img1 = (tonumpy(unnormalize(img1)) * 255.).astype(np.uint8)
desc1 = (desc2RGB(tonumpy(desc1)) * 255.).astype(np.uint8)
kps1 = kps1.detach().cpu().numpy()
# compute the corresponding region
H = H.detach().cpu().numpy()
h1, w1 = img1.shape[:2]
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]
# draw the corresponding region on the image
pts0 = pts0.astype(np.int32)
img0 = cv2.polylines(img0.copy(), [pts0.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
desc0 = cv2.polylines(desc0.copy(), [pts0.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
# desc0 = cv2.warpPerspective(desc0.copy(), H, (w1, h1), flags=cv2.INTER_LINEAR)
pts1 = pts1.astype(np.int32)
img1 = cv2.polylines(img1.copy(), [pts1.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
desc1 = cv2.polylines(desc1.copy(), [pts1.reshape([-1, 2])], True, (255, 0, 0), thickness=5)
# draw correspondences
img = draw_corr(img0, kps0, img1, kps1, num=10)
# desc = draw_corr(desc0, kps0, desc1, kps1, num=10)
desc = draw_corr(desc0, kps1, desc1, kps1, num=10)
return np.concatenate([img, desc], axis=0)
def evaluate(self, scale_pred, scale, image0, image1, index, mask):
"""
scale_pred: [3, H', W']
scale: [1, H, W]
image0: [3, H, W]
image1: [3, H, W]
mask: [1, H, W], boolean tensor
"""
num_cls = scale_pred.shape[0]
scale_pred = torch.argmax(scale_pred, dim=0).long()
scale[scale > 1.5] = 2
scale[scale < 0.75] = 0
scale[(scale >= 0.75) * (scale <= 1.5)] = 1
scale = scale.long()
h, w = scale_pred.shape
h0, w0 = scale.shape[1:]
scale = F.interpolate(scale.view(1, 1, h0, w0).float(), (h, w),
mode='bilinear').long()[0, 0]
mask = F.interpolate(mask.view(1, 1, h0, w0).float(), (h, w),
mode='bilinear').byte()[0, 0]
# correct predictions
correct = (scale_pred == scale) & mask
accuracy = torch.sum(correct).float() / torch.sum(mask).float()
self.accuracies.append(accuracy)
scale_pred = touint8(tonumpy(cls2RGB(scale_pred, num_cls)))
scale = touint8(tonumpy(cls2RGB(scale, num_cls)))
img0 = tensor2RGB(image0)
img1 = tensor2RGB(image1)
img = np.concatenate([img0, img1], axis=1)
import matplotlib.pyplot as plt
_, (ax1, ax2, ax3) = plt.subplots(1, 3)
ax1.imshow(img)
ax2.imshow(scale)
ax3.imshow(scale_pred)
plt.show()
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 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
}