Spaces:
Sleeping
Sleeping
| 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() | |