| | import os |
| | import cv2 |
| | import numpy as np |
| | import os.path as osp |
| | import imageio |
| | from copy import deepcopy |
| |
|
| | import loguru |
| | import torch |
| | import matplotlib.cm as cm |
| | import matplotlib.pyplot as plt |
| |
|
| | from ..loftr import LoFTR, default_cfg |
| | from . import plt_utils |
| | from .plotting import make_matching_figure |
| | from .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 |
| | matcher = LoFTR(config=_default_cfg) |
| | current_dir = os.path.dirname(os.path.abspath(__file__)) |
| | ckpt_path = os.path.join(current_dir, "weights/indoor_ds_new.ckpt") |
| | if not osp.exists(ckpt_path): |
| | raise FileNotFoundError(f"Checkpoint not found at {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} |
| |
|
| | |
| | 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: |
| | |
| | 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]) |
| | 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: |
| | |
| | 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 |
| |
|
