def process(self, frame):
if self.anchor_frame_tensor is None:
print("Please set anchor frame first...")
return None
frame_tensor = frame2tensor(frame, self.device)
pred = self.matching({**self.anchor_data, "image1": frame_tensor})
kpts0 = self.anchor_data["keypoints0"][0].cpu().numpy()
kpts1 = pred["keypoints1"][0].cpu().numpy()
matches = pred["matches0"][0].cpu().numpy()
confidence = pred["matching_scores0"][0].cpu().numpy()
# print("anchor",self.anchor_data.keys())
# print(self.anchor_data['keypoints0'][0].shape)
# print(self.anchor_data['scores0'][0].shape)
# print(self.anchor_data['descriptors0'][0].shape)
# print("pred",pred.keys())
valid = matches > -1
mkpts0 = kpts0[valid]
mkpts1 = kpts1[matches[valid]]
color = cm.jet(confidence[valid])
text = [
"SuperGlue",
"Keypoints: {}:{}".format(len(kpts0), len(kpts1)),
"Matches: {}".format(len(mkpts0)),
]
k_thresh = self.matching.superpoint.config["keypoint_threshold"]
m_thresh = self.matching.superglue.config["match_threshold"]
small_text = [
"Keypoint Threshold: {:.4f}".format(k_thresh),
"Match Threshold: {:.2f}".format(m_thresh),
]
out = make_matching_plot_fast(
self.anchor_frame,
frame,
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
path=None,
show_keypoints=self.show_keypoints,
small_text=small_text,
)
return out, mkpts0, mkpts1
'SuperGlue', 'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)),
'Matches: {}'.format(len(mkpts0))
]
k_thresh = matching.superpoint.config['keypoint_threshold']
m_thresh = matching.superglue.config['match_threshold']
small_text = [
'Keypoint Threshold: {:.4f}'.format(k_thresh),
'Match Threshold: {:.2f}'.format(m_thresh),
'Image Pair: {:06}:{:06}'.format(stem0, stem1),
]
out = make_matching_plot_fast(last_frame,
frame,
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
path=None,
show_keypoints=opt.show_keypoints,
small_text=small_text)
if not opt.no_display:
cv2.imshow('SuperGlue matches', out)
key = chr(cv2.waitKey(1) & 0xFF)
if key == 'q':
vs.cleanup()
print('Exiting (via q) demo_superglue.py')
break
elif key == 'n': # set the current frame as anchor
last_data = {k + '0': pred[k + '1'] for k in keys}
def getKpts(matching, image0_path, image1_path, n=-1):
torch.set_grad_enabled(False)
keys = ['keypoints', 'scores', 'descriptors']
# make image0 first, image 1 second
if image0_path[-5] == '2':
image0_path, image1_path = image1_path, image0_path
print(image0_path)
frame0 = cv2.imread(image0_path, 0)
print(type(frame0), type(mask0))
frame0 = (torch.Tensor(frame0)).to(device)
frame0 = torch.mul(frame0, mask0) # apply mask
frame_tensor0 = (frame0 / 255.).float()[None, None]
last_data = matching.superpoint({'image': frame_tensor0})
last_data = {k + '0': last_data[k] for k in keys}
last_data['image0'] = frame_tensor0
frame1 = cv2.imread(image1_path, 0)
frame1 = (torch.Tensor(frame1)).to(device)
frame1 = torch.mul(frame1, mask1) # apply mask
frame1 = torch.rot90(frame1) # rotate image to make images 0,1 similar
frame_tensor1 = (frame1 / 255.).float()[None, None]
pred = matching({**last_data, 'image1': frame_tensor1})
kpts0 = last_data['keypoints0'][0]
kpts1 = pred['keypoints1'][0]
matches = pred['matches0'][0]
confidence = pred['matching_scores0'][0]
# valid_data
valid = matches > -1
valid_kpts0 = kpts0[valid].to(device)
valid_kpts1 = kpts1[matches[valid]].to(device)
valid_conf = confidence[valid].to(device)
# print('ind shape before deletion: ', valid_conf.shape)
# img_center = torch.Tensor([pic_half_size,pic_half_size])
if len(valid_conf.size()) > 0:
# delete out-of-circle points
center = torch.Tensor([pic_half_size, pic_half_size]).to(device)
len0 = torch.norm(valid_kpts0 - center, dim=1)
len1 = torch.norm(valid_kpts1 - center, dim=1)
indices0 = (len0 < pic_half_size).nonzero()
indices1 = (len1 < pic_half_size).nonzero()
indices = torch.unique(torch.cat((indices0, indices1)))
valid_kpts0 = torch.squeeze(valid_kpts0[indices])
valid_kpts1 = torch.squeeze(valid_kpts1[indices])
valid_conf = torch.squeeze(valid_conf[indices])
if len(valid_conf.size()) > 0:
frame0_brightness = torch.zeros(valid_conf.shape[0])
frame1_brightness = torch.zeros(valid_conf.shape[0])
for k in range(valid_conf.shape[0]):
frame0_brightness[k] = frame0[int(valid_kpts0[k, 0]), int(valid_kpts0[k, 1])] / 255
frame1_brightness[k] = frame1[int(valid_kpts1[k, 0]), int(valid_kpts1[k, 1])] / 255
if n != -1:
if valid_conf.shape[0] < n:
valid_kpts0 = torch.zeros((n, 2))
valid_kpts1 = torch.zeros((n, 2))
valid_conf = torch.zeros(n)
frame0_brightness = torch.zeros(n)
frame1_brightness = torch.zeros(n)
not_skip = 0
# print('here0')
else:
valid_conf, arr1inds = torch.sort(valid_conf, descending=True)
valid_kpts0 = valid_kpts0[arr1inds, :]
valid_kpts1 = valid_kpts1[arr1inds, :]
frame0_brightness = frame0_brightness[arr1inds]
frame1_brightness = frame1_brightness[arr1inds]
valid_kpts0 = valid_kpts0[0:n]
valid_kpts1 = valid_kpts1[0:n]
valid_conf = valid_conf[0:n]
frame0_brightness = frame0_brightness[0:n]
frame1_brightness = frame1_brightness[0:n]
not_skip = 1
# print('here1')
# if n==-1, keep all points
else:
not_skip = 1
# print('here2')
else:
valid_kpts0 = torch.zeros((1, 2))
valid_kpts1 = torch.zeros((1, 2))
valid_conf = torch.zeros(2)
frame0_brightness = torch.zeros(2)
frame1_brightness = torch.zeros(2)
not_skip = 0
# print('here3')
# save image
if not_skip:
valid_kpts0_pic = valid_kpts0.cpu().numpy()
valid_kpts1_pic = valid_kpts1.cpu().numpy()
valid_conf_pic = valid_conf.cpu().numpy()
frame0_pic = frame0.cpu().numpy()
frame1_pic = frame1.cpu().numpy()
kpts0_pic = kpts0.cpu().numpy()
kpts1_pic = kpts1.cpu().numpy()
color = cm.jet(valid_conf_pic)
text = [
'SuperGlue',
'Keypoints: {}:{}'.format(len(kpts0_pic), len(kpts1_pic)),
'Matches: {}'.format(len(valid_kpts0_pic))
]
k_thresh = matching.superpoint.config['keypoint_threshold']
m_thresh = matching.superglue.config['match_threshold']
small_text = [
'Keypoint Threshold: {:.4f}'.format(k_thresh),
'Match Threshold: {:.2f}'.format(m_thresh),
]
out = make_matching_plot_fast(
frame0_pic, frame1_pic, kpts0_pic, kpts1_pic, valid_kpts0_pic, valid_kpts1_pic, color, text,
path=None, show_keypoints=True, small_text=small_text)
else:
out = np.zeros((1920, 3850, 3))
return valid_kpts0, valid_kpts1, valid_conf, frame0_brightness, frame1_brightness, not_skip, out
confidence = pred['matching_scores0'][0].cpu().numpy()
timer.update('forward')
valid = matches > -1
mkpts0 = kpts0[valid]
mkpts1 = kpts1[matches[valid]]
color = cm.jet(confidence[valid])
text = [
'SuperGlue', 'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)),
'Matches: {}'.format(len(mkpts0))
]
out = make_matching_plot_fast(last_frame,
frame,
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
add_keypoints=opt.show_keypoints)
if not opt.no_display:
ds = opt.display_scale
out = cv2.resize(out,
(int(out.shape[1] * ds), int(out.shape[0] * ds)))
cv2.imshow('SuperGlue matches', out)
key = chr(cv2.waitKey(1) & 0xFF)
if key == 'q':
print('Exiting (via q) demo_superglue.py')
break
elif key == 'n': # set the current frame as reference
]
k_thresh = matching.superpoint.config['keypoint_threshold']
m_thresh = matching.superglue.config['match_threshold']
small_text = [
'Keypoint Threshold: {:.4f}'.format(k_thresh),
'Match Threshold: {:.2f}'.format(m_thresh),
'Image Pair: {:06}:{:06}'.format(stem0, stem1),
]
if opt.regionnetvlad:
out = make_matching_plot_fast(cv2.cvtColor(last_frame,
cv2.COLOR_RGB2GRAY),
cv2.cvtColor(frame,
cv2.COLOR_RGB2GRAY),
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
path=None,
show_keypoints=opt.show_keypoints,
small_text=small_text)
else:
out = make_matching_plot_fast(last_frame,
frame,
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
