Spaces:
Running
Running
| from argparse import Namespace | |
| import os, sys | |
| import torch | |
| import cv2 | |
| from pathlib import Path | |
| from .base import Viz | |
| from src.utils.metrics import compute_symmetrical_epipolar_errors, compute_pose_errors | |
| patch2pix_path = Path(__file__).parent / "../../third_party/patch2pix" | |
| sys.path.append(str(patch2pix_path)) | |
| from third_party.patch2pix.utils.eval.model_helper import load_model, estimate_matches | |
| class VizPatch2Pix(Viz): | |
| def __init__(self, args): | |
| super().__init__() | |
| if type(args) == dict: | |
| args = Namespace(**args) | |
| self.imsize = args.imsize | |
| self.match_threshold = args.match_threshold | |
| self.ksize = args.ksize | |
| self.model = load_model(args.ckpt, method="patch2pix") | |
| self.name = "Patch2Pix" | |
| print(f"Initialize {self.name} with image size {self.imsize}") | |
| def match_and_draw( | |
| self, | |
| data_dict, | |
| root_dir=None, | |
| ground_truth=False, | |
| measure_time=False, | |
| viz_matches=True, | |
| ): | |
| img_name0, img_name1 = list(zip(*data_dict["pair_names"]))[0] | |
| path_img0 = os.path.join(root_dir, img_name0) | |
| path_img1 = os.path.join(root_dir, img_name1) | |
| img0, img1 = cv2.imread(path_img0), cv2.imread(path_img1) | |
| return_m_upscale = True | |
| if str(data_dict["dataset_name"][0]).lower() == "scannet": | |
| # self.imsize = 640 | |
| img0 = cv2.resize(img0, tuple(self.imsize)) # (640, 480)) | |
| img1 = cv2.resize(img1, tuple(self.imsize)) # (640, 480)) | |
| return_m_upscale = False | |
| outputs = estimate_matches( | |
| self.model, | |
| path_img0, | |
| path_img1, | |
| ksize=self.ksize, | |
| io_thres=self.match_threshold, | |
| eval_type="fine", | |
| imsize=self.imsize, | |
| return_upscale=return_m_upscale, | |
| measure_time=measure_time, | |
| ) | |
| if measure_time: | |
| self.time_stats.append(outputs[-1]) | |
| matches, mconf = outputs[0], outputs[1] | |
| kpts0 = matches[:, :2] | |
| kpts1 = matches[:, 2:4] | |
| 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: | |
| data_dict["mkpts0_f"] = ( | |
| torch.from_numpy(matches[:, :2]).float().to(self.device) | |
| ) | |
| data_dict["mkpts1_f"] = ( | |
| torch.from_numpy(matches[:, 2:4]).float().to(self.device) | |
| ) | |
| data_dict["m_bids"] = torch.zeros( | |
| matches.shape[0], device=self.device, dtype=torch.float32 | |
| ) | |
| 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, | |
| ) | |
| 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 - mconf | |
| self.draw_matches( | |
| kpts0, | |
| kpts1, | |
| img0, | |
| img1, | |
| m_conf, | |
| path=path_to_save_matches, | |
| conf_thr=0.4, | |
| ) | |