import matplotlib.pyplot as plt import warnings import numpy as np import cv2 import os import os.path as osp import imageio from copy import deepcopy import loguru import torch from oee.models.loftr import LoFTR, default_cfg import matplotlib.cm as cm from oee.utils import plt_utils from oee.utils.plotting import make_matching_figure from oee.utils.utils3d import rect_to_img, canonical_to_camera, calc_pose class ElevEstHelper: _feature_matcher = None @classmethod def get_feature_matcher(cls): if cls._feature_matcher is None: loguru.logger.info("Loading feature matcher...") _default_cfg = deepcopy(default_cfg) _default_cfg['coarse']['temp_bug_fix'] = True # set to False when using the old ckpt matcher = LoFTR(config=_default_cfg) ckpt_path = "weights/indoor_ds_new.ckpt" if not osp.exists(ckpt_path): loguru.logger.info("Downloading feature matcher...") os.makedirs("weights", exist_ok=True) import gdown gdown.cached_download(url="https://drive.google.com/uc?id=19s3QvcCWQ6g-N1PrYlDCg-2mOJZ3kkgS", path=ckpt_path) matcher.load_state_dict(torch.load(ckpt_path)['state_dict']) matcher = matcher.eval().cuda() cls._feature_matcher = matcher return cls._feature_matcher def mask_out_bkgd(img_path, dbg=False): img = imageio.imread_v2(img_path) if img.shape[-1] == 4: fg_mask = img[:, :, :3] else: loguru.logger.info("Image has no alpha channel, using thresholding to mask out background") fg_mask = ~(img > 245).all(axis=-1) if dbg: plt.imshow(plt_utils.vis_mask(img, fg_mask.astype(np.uint8), color=[0, 255, 0])) plt.show() return fg_mask def get_feature_matching(img_paths, dbg=False): assert len(img_paths) == 4 matcher = ElevEstHelper.get_feature_matcher() feature_matching = {} masks = [] for i in range(4): mask = mask_out_bkgd(img_paths[i], dbg=dbg) masks.append(mask) for i in range(0, 4): for j in range(i + 1, 4): img0_pth = img_paths[i] img1_pth = img_paths[j] mask0 = masks[i] mask1 = masks[j] img0_raw = cv2.imread(img0_pth, cv2.IMREAD_GRAYSCALE) img1_raw = cv2.imread(img1_pth, cv2.IMREAD_GRAYSCALE) original_shape = img0_raw.shape img0_raw_resized = cv2.resize(img0_raw, (480, 480)) img1_raw_resized = cv2.resize(img1_raw, (480, 480)) img0 = torch.from_numpy(img0_raw_resized)[None][None].cuda() / 255. img1 = torch.from_numpy(img1_raw_resized)[None][None].cuda() / 255. batch = {'image0': img0, 'image1': img1} # Inference with LoFTR and get prediction with torch.no_grad(): matcher(batch) mkpts0 = batch['mkpts0_f'].cpu().numpy() mkpts1 = batch['mkpts1_f'].cpu().numpy() mconf = batch['mconf'].cpu().numpy() mkpts0[:, 0] = mkpts0[:, 0] * original_shape[1] / 480 mkpts0[:, 1] = mkpts0[:, 1] * original_shape[0] / 480 mkpts1[:, 0] = mkpts1[:, 0] * original_shape[1] / 480 mkpts1[:, 1] = mkpts1[:, 1] * original_shape[0] / 480 keep0 = mask0[mkpts0[:, 1].astype(int), mkpts1[:, 0].astype(int)] keep1 = mask1[mkpts1[:, 1].astype(int), mkpts1[:, 0].astype(int)] keep = np.logical_and(keep0, keep1) mkpts0 = mkpts0[keep] mkpts1 = mkpts1[keep] mconf = mconf[keep] if dbg: # Draw visualization color = cm.jet(mconf) text = [ 'LoFTR', 'Matches: {}'.format(len(mkpts0)), ] fig = make_matching_figure(img0_raw, img1_raw, mkpts0, mkpts1, color, text=text) fig.show() feature_matching[f"{i}_{j}"] = np.concatenate([mkpts0, mkpts1, mconf[:, None]], axis=1) return feature_matching def gen_pose_hypothesis(center_elevation): elevations = np.radians( [center_elevation, center_elevation - 10, center_elevation + 10, center_elevation, center_elevation]) # 45~120 azimuths = np.radians([30, 30, 30, 20, 40]) input_poses = calc_pose(elevations, azimuths, len(azimuths)) input_poses = input_poses[1:] input_poses[..., 1] *= -1 input_poses[..., 2] *= -1 return input_poses def ba_error_general(K, matches, poses): projmat0 = K @ poses[0].inverse()[:3, :4] projmat1 = K @ poses[1].inverse()[:3, :4] match_01 = matches[0] pts0 = match_01[:, :2] pts1 = match_01[:, 2:4] Xref = cv2.triangulatePoints(projmat0.cpu().numpy(), projmat1.cpu().numpy(), pts0.cpu().numpy().T, pts1.cpu().numpy().T) Xref = Xref[:3] / Xref[3:] Xref = Xref.T Xref = torch.from_numpy(Xref).cuda().float() reproj_error = 0 for match, cp in zip(matches[1:], poses[2:]): dist = (torch.norm(match_01[:, :2][:, None, :] - match[:, :2][None, :, :], dim=-1)) if dist.numel() > 0: # print("dist.shape", dist.shape) m0to2_index = dist.argmin(1) keep = dist[torch.arange(match_01.shape[0]), m0to2_index] < 1 if keep.sum() > 0: xref_in2 = rect_to_img(K, canonical_to_camera(Xref, cp.inverse())) reproj_error2 = torch.norm(match[m0to2_index][keep][:, 2:4] - xref_in2[keep], dim=-1) conf02 = match[m0to2_index][keep][:, -1] reproj_error += (reproj_error2 * conf02).sum() / (conf02.sum()) return reproj_error def find_optim_elev(elevs, nimgs, matches, K, dbg=False): errs = [] for elev in elevs: err = 0 cam_poses = gen_pose_hypothesis(elev) for start in range(nimgs - 1): batch_matches, batch_poses = [], [] for i in range(start, nimgs + start): ci = i % nimgs batch_poses.append(cam_poses[ci]) for j in range(nimgs - 1): key = f"{start}_{(start + j + 1) % nimgs}" match = matches[key] batch_matches.append(match) err += ba_error_general(K, batch_matches, batch_poses) errs.append(err) errs = torch.tensor(errs) if dbg: plt.plot(elevs, errs) plt.show() optim_elev = elevs[torch.argmin(errs)].item() return optim_elev def get_elev_est(feature_matching, min_elev=30, max_elev=150, K=None, dbg=False): flag = True matches = {} for i in range(4): for j in range(i + 1, 4): match_ij = feature_matching[f"{i}_{j}"] if len(match_ij) == 0: flag = False match_ji = np.concatenate([match_ij[:, 2:4], match_ij[:, 0:2], match_ij[:, 4:5]], axis=1) matches[f"{i}_{j}"] = torch.from_numpy(match_ij).float().cuda() matches[f"{j}_{i}"] = torch.from_numpy(match_ji).float().cuda() if not flag: loguru.logger.info("0 matches, could not estimate elevation") return None interval = 10 elevs = np.arange(min_elev, max_elev, interval) optim_elev1 = find_optim_elev(elevs, 4, matches, K) elevs = np.arange(optim_elev1 - 10, optim_elev1 + 10, 1) optim_elev2 = find_optim_elev(elevs, 4, matches, K) return optim_elev2 def elev_est_api(img_paths, min_elev=30, max_elev=150, K=None, dbg=False): feature_matching = get_feature_matching(img_paths, dbg=dbg) if K is None: loguru.logger.warning("K is not provided, using default K") K = np.array([[280.0, 0, 128.0], [0, 280.0, 128.0], [0, 0, 1]]) K = torch.from_numpy(K).cuda().float() elev = get_elev_est(feature_matching, min_elev, max_elev, K, dbg=dbg) return elev