from argparse import Namespace import os import torch import cv2 from time import time from pathlib import Path import matplotlib.cm as cm import numpy as np from src.models.topic_fm import TopicFM from src import get_model_cfg from .base import Viz from src.utils.metrics import compute_symmetrical_epipolar_errors, compute_pose_errors from src.utils.plotting import draw_topics, draw_topicfm_demo, error_colormap class VizTopicFM(Viz): def __init__(self, args): super().__init__() if type(args) == dict: args = Namespace(**args) self.match_threshold = args.match_threshold self.n_sampling_topics = args.n_sampling_topics self.show_n_topics = args.show_n_topics # Load model conf = dict(get_model_cfg()) conf['match_coarse']['thr'] = self.match_threshold conf['coarse']['n_samples'] = self.n_sampling_topics print("model config: ", conf) self.model = TopicFM(config=conf) ckpt_dict = torch.load(args.ckpt) self.model.load_state_dict(ckpt_dict['state_dict']) self.model = self.model.eval().to(self.device) # Name the method # self.ckpt_name = args.ckpt.split('/')[-1].split('.')[0] self.name = 'TopicFM' print(f'Initialize {self.name}') def match_and_draw(self, data_dict, root_dir=None, ground_truth=False, measure_time=False, viz_matches=True): if measure_time: torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() self.model(data_dict) if measure_time: torch.cuda.synchronize() end.record() torch.cuda.synchronize() self.time_stats.append(start.elapsed_time(end)) kpts0 = data_dict['mkpts0_f'].cpu().numpy() kpts1 = data_dict['mkpts1_f'].cpu().numpy() img_name0, img_name1 = list(zip(*data_dict['pair_names']))[0] img0 = cv2.imread(os.path.join(root_dir, img_name0)) img1 = cv2.imread(os.path.join(root_dir, img_name1)) if str(data_dict["dataset_name"][0]).lower() == 'scannet': img0 = cv2.resize(img0, (640, 480)) img1 = cv2.resize(img1, (640, 480)) if viz_matches: saved_name = "_".join([img_name0.split('/')[-1].split('.')[0], img_name1.split('/')[-1].split('.')[0]]) folder_matches = os.path.join(root_dir, "{}_viz_matches".format(self.name)) if not os.path.exists(folder_matches): os.makedirs(folder_matches) path_to_save_matches = os.path.join(folder_matches, "{}.png".format(saved_name)) if ground_truth: compute_symmetrical_epipolar_errors(data_dict) # compute epi_errs for each match compute_pose_errors(data_dict) # compute R_errs, t_errs, pose_errs for each pair epi_errors = data_dict['epi_errs'].cpu().numpy() R_errors, t_errors = data_dict['R_errs'][0], data_dict['t_errs'][0] self.draw_matches(kpts0, kpts1, img0, img1, epi_errors, path=path_to_save_matches, R_errs=R_errors, t_errs=t_errors) # compute evaluation metrics rel_pair_names = list(zip(*data_dict['pair_names'])) bs = data_dict['image0'].size(0) metrics = { # to filter duplicate pairs caused by DistributedSampler 'identifiers': ['#'.join(rel_pair_names[b]) for b in range(bs)], 'epi_errs': [data_dict['epi_errs'][data_dict['m_bids'] == b].cpu().numpy() for b in range(bs)], 'R_errs': data_dict['R_errs'], 't_errs': data_dict['t_errs'], 'inliers': data_dict['inliers']} self.eval_stats.append({'metrics': metrics}) else: m_conf = 1 - data_dict["mconf"].cpu().numpy() self.draw_matches(kpts0, kpts1, img0, img1, m_conf, path=path_to_save_matches, conf_thr=0.4) if self.show_n_topics > 0: folder_topics = os.path.join(root_dir, "{}_viz_topics".format(self.name)) if not os.path.exists(folder_topics): os.makedirs(folder_topics) draw_topics(data_dict, img0, img1, saved_folder=folder_topics, show_n_topics=self.show_n_topics, saved_name=saved_name) def run_demo(self, dataloader, writer=None, output_dir=None, no_display=False, skip_frames=1): data_dict = next(dataloader) frame_id = 0 last_image_id = 0 img0 = np.array(cv2.imread(str(data_dict["img_path"][0])), dtype=np.float32) / 255 frame_tensor = data_dict["img"].to(self.device) pair_data = {'image0': frame_tensor} last_frame = cv2.resize(img0, (frame_tensor.shape[-1], frame_tensor.shape[-2]), cv2.INTER_LINEAR) if output_dir is not None: print('==> Will write outputs to {}'.format(output_dir)) Path(output_dir).mkdir(exist_ok=True) # Create a window to display the demo. if not no_display: window_name = 'Topic-assisted Feature Matching' cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.resizeWindow(window_name, (640 * 2, 480 * 2)) else: print('Skipping visualization, will not show a GUI.') # Print the keyboard help menu. print('==> Keyboard control:\n' '\tn: select the current frame as the reference image (left)\n' '\tq: quit') # vis_range = [kwargs["bottom_k"], kwargs["top_k"]] while True: frame_id += 1 if frame_id == len(dataloader): print('Finished demo_loftr.py') break data_dict = next(dataloader) if frame_id % skip_frames != 0: # print("Skipping frame.") continue stem0, stem1 = last_image_id, data_dict["id"][0].item() - 1 frame = np.array(cv2.imread(str(data_dict["img_path"][0])), dtype=np.float32) / 255 frame_tensor = data_dict["img"].to(self.device) frame = cv2.resize(frame, (frame_tensor.shape[-1], frame_tensor.shape[-2]), interpolation=cv2.INTER_LINEAR) pair_data = {**pair_data, 'image1': frame_tensor} self.model(pair_data) total_n_matches = len(pair_data['mkpts0_f']) mkpts0 = pair_data['mkpts0_f'].cpu().numpy() # [vis_range[0]:vis_range[1]] mkpts1 = pair_data['mkpts1_f'].cpu().numpy() # [vis_range[0]:vis_range[1]] mconf = pair_data['mconf'].cpu().numpy() # [vis_range[0]:vis_range[1]] # Normalize confidence. if len(mconf) > 0: mconf = 1 - mconf # alpha = 0 # color = cm.jet(mconf, alpha=alpha) color = error_colormap(mconf, thr=0.4, alpha=0.1) text = [ f'Topics', '#Matches: {}'.format(total_n_matches), ] out = draw_topicfm_demo(pair_data, last_frame, frame, mkpts0, mkpts1, color, text, show_n_topics=4, path=None) if not no_display: if writer is not None: writer.write(out) cv2.imshow('TopicFM Matches', out) key = chr(cv2.waitKey(10) & 0xFF) if key == 'q': if writer is not None: writer.release() print('Exiting...') break elif key == 'n': pair_data['image0'] = frame_tensor last_frame = frame last_image_id = (data_dict["id"][0].item() - 1) frame_id_left = frame_id elif output_dir is not None: stem = 'matches_{:06}_{:06}'.format(stem0, stem1) out_file = str(Path(output_dir, stem + '.png')) print('\nWriting image to {}'.format(out_file)) cv2.imwrite(out_file, out) else: raise ValueError("output_dir is required when no display is given.") cv2.destroyAllWindows() if writer is not None: writer.release()