|
|
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': |
|
|
|
|
|
img0 = cv2.resize(img0, tuple(self.imsize)) |
|
|
img1 = cv2.resize(img1, tuple(self.imsize)) |
|
|
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_pose_errors(data_dict) |
|
|
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 = { |
|
|
|
|
|
'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) |
|
|
|
