vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/datasets/common_data_pair.py

@@ -0,0 +1,214 @@
+import os.path as osp
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from loguru import logger
+from PIL import Image
+
+from src.utils.dataset import read_megadepth_gray
+
+class CommonDataset(Dataset):
+    def __init__(self,
+                 root_dir,
+                 npz_path,
+                 mode='train',
+                 min_overlap_score=0.4,
+                 img_resize=None,
+                 df=None,
+                 img_padding=False,
+                 depth_padding=False,
+                 augment_fn=None,
+                 testNpairs=300,
+                 fp16=False,
+                 fix_bias=False,
+                 sample_ratio=1.0,
+                 **kwargs):
+        super().__init__()
+        self.root_dir = root_dir
+        self.mode = mode
+        self.scene_id = npz_path.split('.')[0]
+        self.sample_ratio = sample_ratio
+
+        # prepare scene_info and pair_info
+        if mode == 'test' and min_overlap_score > 0:
+            logger.warning("You are using `min_overlap_score`!=0 in test mode. Set to 0.")
+            min_overlap_score = -3.0
+        self.scene_info = np.load(npz_path, allow_pickle=True)
+        if mode == 'test':
+            self.pair_infos = self.scene_info['pair_infos'][:testNpairs].copy()
+        else:
+            self.pair_infos = self.scene_info['pair_infos'].copy()
+
+        # parameters for image resizing, padding and depthmap padding
+        if mode == 'train':
+            assert img_resize is not None and depth_padding
+        self.img_resize = img_resize
+        self.df = df
+        self.img_padding = img_padding
+
+        # for training LoFTR
+        self.augment_fn = augment_fn if mode == 'train' else None
+        self.coarse_scale = getattr(kwargs, 'coarse_scale', 0.125)
+        self.load_origin_rgb = kwargs["load_origin_rgb"]
+        self.read_gray = kwargs["read_gray"]
+        self.normalize_img = kwargs["normalize_img"]
+        self.resize_by_stretch = kwargs["resize_by_stretch"]
+        depth_max_size = 4000 if 'MTV_cross_modal_data' not in npz_path else 6000
+        self.depth_max_size = depth_max_size if depth_padding else 2000  # the upperbound of depthmaps size in megadepth.
+
+        self.dataset_name = self.scene_info['dataset_name'] if "dataset_name" in self.scene_info else npz_path.split(root_dir)[1].split('/')[1]
+        self.gt_matches = self.scene_info['gt_matches'] if 'gt_matches' in self.scene_info else None # sparse matches produced by teacher models, used for training
+        self.gt_matches_padding_n = kwargs["gt_matches_padding_n"]
+        self.gt_2D_warp = self.scene_info['gt_2D_transforms'] if "gt_2D_transforms" in self.scene_info else None
+        self.gt_2D_matches = self.scene_info['gt_2D_matches'] if "gt_2D_matches" in self.scene_info else None # Used for eval
+        self.intrins = self.scene_info['intrinsics'] if 'intrinsics' in self.scene_info else None
+        self.poses = self.scene_info['poses'] if 'poses' in self.scene_info else None
+        
+        self.fp16 = fp16
+        self.fix_bias = fix_bias
+        if self.fix_bias:
+            self.df = 1          
+  
+    def __len__(self):
+        return len(self.pair_infos)
+
+    def __getitem__(self, idx):
+        if isinstance(self.pair_infos[idx], np.ndarray):
+            idx0, idx1 = self.pair_infos[idx][0], self.pair_infos[idx][1]
+            img_path0, img_path1 = self.scene_info['image_paths'][idx0][0], self.scene_info['image_paths'][idx1][1]
+            K_0 = torch.zeros((3,3), dtype=torch.float) if self.intrins is None else torch.from_numpy(self.intrins[idx0][0]).float()
+            K_1 = torch.zeros((3,3), dtype=torch.float) if self.intrins is None else torch.from_numpy(self.intrins[idx1][1]).float()
+
+        else:
+            if len(self.pair_infos[idx]) == 3:
+                (idx0, idx1), overlap_score, central_matches = self.pair_infos[idx]
+            elif len(self.pair_infos[idx]) == 2:
+                (idx0, idx1), overlap_score = self.pair_infos[idx]
+            else:
+                raise NotImplementedError
+
+            img_path0, img_path1 = self.scene_info['image_paths'][idx0], self.scene_info['image_paths'][idx1]
+            K_0 = torch.zeros((3,3), dtype=torch.float) if self.intrins is None else torch.from_numpy(self.intrins[idx0]).float()
+            K_1 = torch.zeros((3,3), dtype=torch.float) if self.intrins is None else torch.from_numpy(self.intrins[idx1]).float()
+
+        # read grayscale image and mask. (1, h, w) and (h, w)
+        img_name0 = osp.join(self.root_dir, self.dataset_name, img_path0)
+        img_name1 = osp.join(self.root_dir, self.dataset_name, img_path1) # Often transformed image based on img0, e.g., depth estimation or Diffusion
+        # Note: should be pixel aligned with img0
+
+        image0, mask0, scale0, origin_img_size0 = read_megadepth_gray(
+            img_name0, self.img_resize, self.df, self.img_padding, None, read_gray=self.read_gray, normalize_img=self.normalize_img, resize_by_stretch=self.resize_by_stretch)
+            # np.random.choice([self.augment_fn, None], p=[0.5, 0.5]))
+        image1, mask1, scale1, origin_img_size1 = read_megadepth_gray(
+            img_name1, self.img_resize, self.df, self.img_padding, None, read_gray=self.read_gray, normalize_img=self.normalize_img, resize_by_stretch=self.resize_by_stretch)
+
+        if self.gt_2D_warp is not None:
+            gt_warp = np.concatenate([self.gt_2D_warp[idx], [[0,0,1]]]) # 3 * 3
+        else:
+            gt_warp = np.zeros((3, 3))
+
+        depth0 = depth1 = torch.zeros([self.depth_max_size, self.depth_max_size], dtype=torch.float)
+
+        homo_mask0 = torch.zeros((1, image0.shape[-2], image0.shape[-1]))
+        homo_mask1 = torch.zeros((1, image1.shape[-2], image1.shape[-1]))
+        gt_matches = torch.zeros((self.gt_matches_padding_n, 4), dtype=torch.float)
+
+        if self.poses is None:
+            T_0to1 = T_1to0 = torch.zeros((4,4), dtype=torch.float)  # (4, 4)
+        else:
+            # read and compute relative poses
+            T0 = self.poses[idx0]
+            T1 = self.poses[idx1]
+            T_0to1 = torch.tensor(np.matmul(T1, np.linalg.inv(T0)), dtype=torch.float)[:4, :4]  # (4, 4)
+            T_1to0 = T_0to1.inverse()
+
+        if self.fp16:
+            data = {
+                'image0': image0.half(),  # (1, h, w)
+                'depth0': depth0.half(),  # (h, w)
+                'image1': image1.half(),
+                'depth1': depth1.half(),
+                'T_0to1': T_0to1,  # (4, 4)
+                'T_1to0': T_1to0,
+                'K0': K_0,  # (3, 3)
+                'K1': K_1,
+                'homo_mask0': homo_mask0,
+                'homo_mask1': homo_mask1,
+                'gt_matches': gt_matches,
+                'gt_matches_mask': torch.zeros((1,), dtype=torch.bool),
+                'homography': torch.from_numpy(gt_warp.astype(np.float32)),
+                'norm_pixel_mat': torch.zeros((3,3), dtype=torch.float),
+                'homo_sample_normed': torch.zeros((3,3), dtype=torch.float),
+                'origin_img_size0': origin_img_size0,
+                'origin_img_size1': origin_img_size1,
+                'scale0': scale0.half(),  # [scale_w, scale_h]
+                'scale1': scale1.half(),
+                'dataset_name': 'MegaDepth',
+                'scene_id': self.scene_id,
+                'pair_id': idx,
+                'pair_names': (img_path0, img_path1),
+            }
+        else:
+            data = {
+                'image0': image0,  # (1, h, w)
+                'depth0': depth0,  # (h, w)
+                'image1': image1,
+                'depth1': depth1,
+                'T_0to1': T_0to1,  # (4, 4)
+                'T_1to0': T_1to0,
+                'K0': K_0,  # (3, 3)
+                'K1': K_1,
+                'homo_mask0': homo_mask0,
+                'homo_mask1': homo_mask1,
+                'homography': torch.from_numpy(gt_warp.astype(np.float32)),
+                'norm_pixel_mat': torch.zeros((3,3), dtype=torch.float),
+                'homo_sample_normed': torch.zeros((3,3), dtype=torch.float),
+                'gt_matches': gt_matches,
+                'gt_matches_mask': torch.zeros((1,), dtype=torch.bool),
+                'origin_img_size0': origin_img_size0, # H W
+                'origin_img_size1': origin_img_size1,
+                'scale0': scale0,  # [scale_w, scale_h]
+                'scale1': scale1,
+                'dataset_name': 'MegaDepth',
+                'scene_id': self.scene_id,
+                'pair_id': idx,
+                'pair_names': (img_path0, img_path1),
+                'rel_pair_names': (img_path0, img_path1)
+            }
+        
+        if self.gt_2D_matches is not None:
+            data.update({'gt_2D_matches': torch.from_numpy(self.gt_2D_matches[idx]).to(torch.float)}) # N * 4
+
+        if self.gt_matches is not None:
+            gt_matches_ = self.gt_matches[idx]
+            if isinstance(gt_matches_, str):
+                gt_matches_ = np.load(osp.join(self.root_dir, self.dataset_name, gt_matches_), allow_pickle=True)
+            gt_matches_ = torch.from_numpy(gt_matches_).to(torch.float) # N * 4: mkpts0, mkpts1
+            # Warp mkpts1 by sampled homo:
+            num = min(len(gt_matches_), self.gt_matches_padding_n)
+            gt_matches[:num] = gt_matches_[:num]
+
+            data.update({"gt_matches": gt_matches, 'gt_matches_mask': torch.ones((1,), dtype=torch.bool), 'norm_pixel_mat': torch.zeros((3,3), dtype=torch.float), "homo_sample_normed": torch.zeros((3,3), dtype=torch.float)})
+
+        if mask0 is not None:  # img_padding is True
+            if self.coarse_scale:
+                if self.fix_bias:
+                    [ts_mask_0, ts_mask_1] = F.interpolate(torch.stack([mask0, mask1], dim=0)[None].float(),
+                                                            size=((image0.shape[1]-1)//8+1, (image0.shape[2]-1)//8+1),
+                                                            mode='nearest',
+                                                            recompute_scale_factor=False)[0].bool()
+                else:
+                    [ts_mask_0, ts_mask_1] = F.interpolate(torch.stack([mask0, mask1], dim=0)[None].float(),
+                                                            scale_factor=self.coarse_scale,
+                                                            mode='nearest',
+                                                            recompute_scale_factor=False)[0].bool()
+            if self.fp16:
+                data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1})
+            else:
+                data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1})
+        
+        if self.load_origin_rgb:
+            data.update({"image0_rgb_origin": torch.from_numpy(np.array(Image.open(img_name0).convert("RGB"))).permute(2,0,1) / 255., "image1_rgb_origin": torch.from_numpy(np.array(Image.open(img_name1).convert("RGB"))).permute(2,0,1)/ 255.})
+
+        return data

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/lightning/lightning_loftr.py

@@ -0,0 +1,343 @@
+
+from collections import defaultdict
+import pprint
+from loguru import logger
+from pathlib import Path
+
+import torch
+import numpy as np
+import pytorch_lightning as pl
+from matplotlib import pyplot as plt
+
+from src.loftr import LoFTR
+from src.loftr.utils.supervision import compute_supervision_coarse, compute_supervision_fine, compute_roma_supervision
+from src.optimizers import build_optimizer, build_scheduler
+from src.utils.metrics import (
+    compute_symmetrical_epipolar_errors,
+    compute_pose_errors,
+    compute_homo_corner_warp_errors,
+    compute_homo_match_warp_errors,
+    compute_warp_control_pts_errors,
+    aggregate_metrics
+)
+from src.utils.plotting import make_matching_figures, make_scores_figures
+from src.utils.comm import gather, all_gather
+from src.utils.misc import lower_config, flattenList
+from src.utils.profiler import PassThroughProfiler
+from third_party.ROMA.roma.matchanything_roma_model import MatchAnything_Model
+
+import pynvml
+
+def reparameter(matcher):
+    module = matcher.backbone.layer0
+    if hasattr(module, 'switch_to_deploy'):
+        module.switch_to_deploy()
+        print('m0 switch to deploy ok')
+    for modules in [matcher.backbone.layer1, matcher.backbone.layer2, matcher.backbone.layer3]:
+        for module in modules:
+            if hasattr(module, 'switch_to_deploy'):
+                module.switch_to_deploy()
+                print('backbone switch to deploy ok')
+    for modules in [matcher.fine_preprocess.layer2_outconv2, matcher.fine_preprocess.layer1_outconv2]:
+        for module in modules:
+            if hasattr(module, 'switch_to_deploy'):
+                module.switch_to_deploy()
+                print('fpn switch to deploy ok')
+    return matcher
+
+class PL_LoFTR(pl.LightningModule):
+    def __init__(self, config, pretrained_ckpt=None, profiler=None, dump_dir=None, test_mode=False, baseline_config=None):
+        """
+        TODO:
+            - use the new version of PL logging API.
+        """
+        super().__init__()
+        # Misc
+        self.config = config  # full config
+        _config = lower_config(self.config)
+        self.profiler = profiler or PassThroughProfiler()
+        self.n_vals_plot = max(config.TRAINER.N_VAL_PAIRS_TO_PLOT // config.TRAINER.WORLD_SIZE, 1)
+
+        if config.METHOD == "matchanything_eloftr":
+            self.matcher = LoFTR(config=_config['loftr'], profiler=self.profiler)
+        elif config.METHOD == "matchanything_roma":
+            self.matcher = MatchAnything_Model(config=_config['roma'], test_mode=test_mode)
+        else: 
+            raise NotImplementedError
+
+        if config.FINETUNE.ENABLE and test_mode:
+            # Inference time change model architecture before load pretrained model:
+            raise NotImplementedError
+
+        # Pretrained weights
+        if pretrained_ckpt:
+            if config.METHOD in ["matchanything_eloftr", "matchanything_roma"]:
+                state_dict = torch.load(pretrained_ckpt, map_location='cpu')['state_dict']
+                logger.info(f"Load model from:{self.matcher.load_state_dict(state_dict, strict=False)}")
+            else:
+                raise NotImplementedError
+
+        if self.config.LOFTR.BACKBONE_TYPE == 'RepVGG' and test_mode and (config.METHOD == 'loftr'):
+            module = self.matcher.backbone.layer0
+            if hasattr(module, 'switch_to_deploy'):
+                module.switch_to_deploy()
+                print('m0 switch to deploy ok')
+            for modules in [self.matcher.backbone.layer1, self.matcher.backbone.layer2, self.matcher.backbone.layer3]:
+                for module in modules:
+                    if hasattr(module, 'switch_to_deploy'):
+                        module.switch_to_deploy()
+                        print('m switch to deploy ok')
+        
+        # Testing
+        self.dump_dir = dump_dir
+        self.max_gpu_memory = 0
+        self.GPUID = 0
+        self.warmup = False
+        
+    def gpumem(self, des, gpuid=None):
+        NUM_EXPAND = 1024 * 1024 * 1024
+        gpu_id= self.GPUID if self.GPUID is not None else gpuid
+        handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
+        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        gpu_Used = info.used
+        logger.info(f"GPU {gpu_id} memory used: {gpu_Used / NUM_EXPAND} GB while {des}")
+        # print(des, gpu_Used / NUM_EXPAND)
+        if gpu_Used / NUM_EXPAND > self.max_gpu_memory:
+            self.max_gpu_memory = gpu_Used / NUM_EXPAND
+            logger.info(f"[MAX]GPU {gpu_id} memory used: {gpu_Used / NUM_EXPAND} GB while {des}")
+            print('max_gpu_memory', self.max_gpu_memory)
+
+    def configure_optimizers(self):
+        optimizer = build_optimizer(self, self.config)
+        scheduler = build_scheduler(self.config, optimizer)
+        return [optimizer], [scheduler]
+    
+    def optimizer_step(
+            self, epoch, batch_idx, optimizer, optimizer_idx,
+            optimizer_closure, on_tpu, using_native_amp, using_lbfgs):
+        # learning rate warm up
+        warmup_step = self.config.TRAINER.WARMUP_STEP
+        if self.trainer.global_step < warmup_step:
+            if self.config.TRAINER.WARMUP_TYPE == 'linear':
+                base_lr = self.config.TRAINER.WARMUP_RATIO * self.config.TRAINER.TRUE_LR
+                lr = base_lr + \
+                    (self.trainer.global_step / self.config.TRAINER.WARMUP_STEP) * \
+                    abs(self.config.TRAINER.TRUE_LR - base_lr)
+                for pg in optimizer.param_groups:
+                    pg['lr'] = lr
+            elif self.config.TRAINER.WARMUP_TYPE == 'constant':
+                pass
+            else:
+                raise ValueError(f'Unknown lr warm-up strategy: {self.config.TRAINER.WARMUP_TYPE}')
+
+        # update params
+        if self.config.LOFTR.FP16:
+            optimizer.step(closure=optimizer_closure)
+        else:
+            optimizer.step(closure=optimizer_closure)
+        optimizer.zero_grad()
+    
+    def _trainval_inference(self, batch):
+        with self.profiler.profile("Compute coarse supervision"):
+
+            with torch.autocast(enabled=False, device_type='cuda'):
+                if ("roma" in self.config.METHOD) or ('dkm' in self.config.METHOD):
+                    pass
+                else:
+                    compute_supervision_coarse(batch, self.config)
+        
+        with self.profiler.profile("LoFTR"):
+            with torch.autocast(enabled=self.config.LOFTR.FP16, device_type='cuda'):
+                self.matcher(batch)
+        
+        with self.profiler.profile("Compute fine supervision"):
+            with torch.autocast(enabled=False, device_type='cuda'):
+                if ("roma" in self.config.METHOD) or ('dkm' in self.config.METHOD):
+                    compute_roma_supervision(batch, self.config)
+                else:
+                    compute_supervision_fine(batch, self.config, self.logger)
+            
+        with self.profiler.profile("Compute losses"):
+            pass
+    
+    def _compute_metrics(self, batch):
+        if 'gt_2D_matches' in batch:
+            compute_warp_control_pts_errors(batch, self.config)
+        elif batch['homography'].sum() != 0 and batch['T_0to1'].sum() == 0:
+            compute_homo_match_warp_errors(batch, self.config)  # compute warp_errors for each match
+            compute_homo_corner_warp_errors(batch, self.config) # compute mean corner warp error each pair
+        else:
+            compute_symmetrical_epipolar_errors(batch, self.config)  # compute epi_errs for each match
+            compute_pose_errors(batch, self.config)  # compute R_errs, t_errs, pose_errs for each pair
+
+        rel_pair_names = list(zip(*batch['pair_names']))
+        bs = batch['image0'].size(0)
+        if self.config.LOFTR.FINE.MTD_SPVS:
+            topk = self.config.LOFTR.MATCH_FINE.TOPK
+            metrics = {
+                # to filter duplicate pairs caused by DistributedSampler
+                'identifiers': ['#'.join(rel_pair_names[b]) for b in range(bs)],
+                'epi_errs': [(batch['epi_errs'].reshape(-1,topk))[batch['m_bids'] == b].reshape(-1).cpu().numpy() for b in range(bs)],
+                'R_errs': batch['R_errs'],
+                't_errs': batch['t_errs'],
+                'inliers': batch['inliers'],
+                'num_matches': [batch['mconf'].shape[0]], # batch size = 1 only
+                'percent_inliers': [ batch['inliers'][0].shape[0] / batch['mconf'].shape[0] if batch['mconf'].shape[0]!=0 else 1], # batch size = 1 only
+                }
+        else:
+            metrics = {
+                # to filter duplicate pairs caused by DistributedSampler
+                'identifiers': ['#'.join(rel_pair_names[b]) for b in range(bs)],
+                'epi_errs': [batch['epi_errs'][batch['m_bids'] == b].cpu().numpy() for b in range(bs)],
+                'R_errs': batch['R_errs'],
+                't_errs': batch['t_errs'],
+                'inliers': batch['inliers'],
+                'num_matches': [batch['mconf'].shape[0]], # batch size = 1 only
+                'percent_inliers': [ batch['inliers'][0].shape[0] / batch['mconf'].shape[0] if batch['mconf'].shape[0]!=0 else 1], # batch size = 1 only
+                }
+        ret_dict = {'metrics': metrics}
+        return ret_dict, rel_pair_names
+    
+    def training_step(self, batch, batch_idx):
+        self._trainval_inference(batch)
+        
+        # logging
+        if self.trainer.global_rank == 0 and self.global_step % self.trainer.log_every_n_steps == 0:
+            # scalars
+            for k, v in batch['loss_scalars'].items():
+                self.logger.experiment.add_scalar(f'train/{k}', v, self.global_step)
+
+            # net-params
+            method = 'LOFTR'
+            if self.config[method]['MATCH_COARSE']['MATCH_TYPE']  == 'sinkhorn':
+                self.logger.experiment.add_scalar(
+                    f'skh_bin_score', self.matcher.coarse_matching.bin_score.clone().detach().cpu().data, self.global_step)
+        
+            figures = {}
+            if self.config.TRAINER.ENABLE_PLOTTING:
+                compute_symmetrical_epipolar_errors(batch, self.config)  # compute epi_errs for each match
+                figures = make_matching_figures(batch, self.config, self.config.TRAINER.PLOT_MODE)
+                for k, v in figures.items():
+                    self.logger.experiment.add_figure(f'train_match/{k}', v, self.global_step)
+
+        return {'loss': batch['loss']}
+
+    def training_epoch_end(self, outputs):
+        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
+        if self.trainer.global_rank == 0:
+            self.logger.experiment.add_scalar(
+                'train/avg_loss_on_epoch', avg_loss,
+                global_step=self.current_epoch)
+    
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        self._trainval_inference(batch)
+        
+        ret_dict, _ = self._compute_metrics(batch)
+        
+        val_plot_interval = max(self.trainer.num_val_batches[0] // self.n_vals_plot, 1)
+        figures = {self.config.TRAINER.PLOT_MODE: []}
+        if batch_idx % val_plot_interval == 0:
+            figures = make_matching_figures(batch, self.config, mode=self.config.TRAINER.PLOT_MODE)
+            if self.config.LOFTR.PLOT_SCORES:
+                figs = make_scores_figures(batch, self.config, self.config.TRAINER.PLOT_MODE)
+                figures[self.config.TRAINER.PLOT_MODE] += figs[self.config.TRAINER.PLOT_MODE]
+                del figs
+                
+        return {
+            **ret_dict,
+            'loss_scalars': batch['loss_scalars'],
+            'figures': figures,
+        }
+        
+    def validation_epoch_end(self, outputs):
+        # handle multiple validation sets
+        multi_outputs = [outputs] if not isinstance(outputs[0], (list, tuple)) else outputs
+        multi_val_metrics = defaultdict(list)
+        
+        for valset_idx, outputs in enumerate(multi_outputs):
+            # since pl performs sanity_check at the very begining of the training
+            cur_epoch = self.trainer.current_epoch
+            if not self.trainer.resume_from_checkpoint and self.trainer.running_sanity_check:
+                cur_epoch = -1
+
+            # 1. loss_scalars: dict of list, on cpu
+            _loss_scalars = [o['loss_scalars'] for o in outputs]
+            loss_scalars = {k: flattenList(all_gather([_ls[k] for _ls in _loss_scalars])) for k in _loss_scalars[0]}
+
+            # 2. val metrics: dict of list, numpy
+            _metrics = [o['metrics'] for o in outputs]
+            metrics = {k: flattenList(all_gather(flattenList([_me[k] for _me in _metrics]))) for k in _metrics[0]}
+            # NOTE: all ranks need to `aggregate_merics`, but only log at rank-0 
+            val_metrics_4tb = aggregate_metrics(metrics, self.config.TRAINER.EPI_ERR_THR, self.config.LOFTR.EVAL_TIMES)
+            for thr in [5, 10, 20]:
+                multi_val_metrics[f'auc@{thr}'].append(val_metrics_4tb[f'auc@{thr}'])
+            
+            # 3. figures
+            _figures = [o['figures'] for o in outputs]
+            figures = {k: flattenList(gather(flattenList([_me[k] for _me in _figures]))) for k in _figures[0]}
+
+            # tensorboard records only on rank 0
+            if self.trainer.global_rank == 0:
+                for k, v in loss_scalars.items():
+                    mean_v = torch.stack(v).mean()
+                    self.logger.experiment.add_scalar(f'val_{valset_idx}/avg_{k}', mean_v, global_step=cur_epoch)
+
+                for k, v in val_metrics_4tb.items():
+                    self.logger.experiment.add_scalar(f"metrics_{valset_idx}/{k}", v, global_step=cur_epoch)
+                
+                for k, v in figures.items():
+                    if self.trainer.global_rank == 0:
+                        for plot_idx, fig in enumerate(v):
+                            self.logger.experiment.add_figure(
+                                f'val_match_{valset_idx}/{k}/pair-{plot_idx}', fig, cur_epoch, close=True)
+            plt.close('all')
+
+        for thr in [5, 10, 20]:
+            self.log(f'auc@{thr}', torch.tensor(np.mean(multi_val_metrics[f'auc@{thr}'])))  # ckpt monitors on this
+
+    def test_step(self, batch, batch_idx):
+        if self.warmup:
+            for i in range(50):
+                self.matcher(batch)
+            self.warmup = False
+
+        with torch.autocast(enabled=self.config.LOFTR.FP16, device_type='cuda'):
+            with self.profiler.profile("LoFTR"):
+                self.matcher(batch)
+
+        ret_dict, rel_pair_names = self._compute_metrics(batch)
+        print(ret_dict['metrics']['num_matches'])
+        self.dump_dir = None
+
+        return ret_dict
+
+    def test_epoch_end(self, outputs):
+        print(self.config)
+        print('max GPU memory: ', self.max_gpu_memory)
+        print(self.profiler.summary())
+        # metrics: dict of list, numpy
+        _metrics = [o['metrics'] for o in outputs]
+        metrics = {k: flattenList(gather(flattenList([_me[k] for _me in _metrics]))) for k in _metrics[0]}
+
+        # [{key: [{...}, *#bs]}, *#batch]
+        if self.dump_dir is not None:
+            Path(self.dump_dir).mkdir(parents=True, exist_ok=True)
+            _dumps = flattenList([o['dumps'] for o in outputs])  # [{...}, #bs*#batch]
+            dumps = flattenList(gather(_dumps))  # [{...}, #proc*#bs*#batch]
+            logger.info(f'Prediction and evaluation results will be saved to: {self.dump_dir}')
+
+        if self.trainer.global_rank == 0:
+            NUM_EXPAND = 1024 * 1024 * 1024
+            gpu_id=self.GPUID
+            handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
+            info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+            gpu_Used = info.used 
+            print('pynvml', gpu_Used / NUM_EXPAND)
+            if gpu_Used / NUM_EXPAND > self.max_gpu_memory:
+                self.max_gpu_memory = gpu_Used / NUM_EXPAND
+
+            print(self.profiler.summary())
+            val_metrics_4tb = aggregate_metrics(metrics, self.config.TRAINER.EPI_ERR_THR, self.config.LOFTR.EVAL_TIMES, self.config.TRAINER.THRESHOLDS, method=self.config.TRAINER.AUC_METHOD)
+            logger.info('\n' + pprint.pformat(val_metrics_4tb))
+            if self.dump_dir is not None:
+                np.save(Path(self.dump_dir) / 'LoFTR_pred_eval', dumps)

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/loftr/__init__.py

@@ -0,0 +1 @@
+from .loftr import LoFTR

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/loftr/backbone/__init__.py

@@ -0,0 +1,61 @@
+from .resnet_fpn import ResNetFPN_8_2, ResNetFPN_16_4, ResNetFPN_8_1, ResNetFPN_8_2_align, ResNetFPN_8_1_align, ResNetFPN_8_2_fix, ResNet_8_1_align, VGG_8_1_align, RepVGG_8_1_align, \
+    RepVGGnfpn_8_1_align, RepVGG_8_2_fix, s2dnet_8_1_align
+
+def build_backbone(config):
+    if config['backbone_type'] == 'ResNetFPN':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            if config['resolution'] == (8, 2):
+                return ResNetFPN_8_2(config['resnetfpn'])
+            elif config['resolution'] == (16, 4):
+                return ResNetFPN_16_4(config['resnetfpn'])
+            elif config['resolution'] == (8, 1):
+                return ResNetFPN_8_1(config['resnetfpn'])
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 2):
+                return ResNetFPN_8_2_align(config['resnetfpn'])
+            elif config['resolution'] == (16, 4):
+                return ResNetFPN_16_4(config['resnetfpn'])
+            elif config['resolution'] == (8, 1):
+                return ResNetFPN_8_1_align(config['resnetfpn'])
+    elif config['backbone_type'] == 'ResNetFPNFIX':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            if config['resolution'] == (8, 2):
+                return ResNetFPN_8_2_fix(config['resnetfpn'])
+    elif config['backbone_type'] == 'ResNet':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 1):
+                return ResNet_8_1_align(config['resnetfpn'])
+    elif config['backbone_type'] == 'VGG':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 1):
+                return VGG_8_1_align(config['resnetfpn'])
+    elif config['backbone_type'] == 'RepVGG':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 1):
+                return RepVGG_8_1_align(config['resnetfpn'])
+    elif config['backbone_type'] == 'RepVGGNFPN':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 1):
+                return RepVGGnfpn_8_1_align(config['resnetfpn'])
+    elif config['backbone_type'] == 'RepVGGFPNFIX':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            if config['resolution'] == (8, 2):
+                return RepVGG_8_2_fix(config['resnetfpn'])
+        elif config['align_corner'] is False:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+    elif config['backbone_type'] == 's2dnet':
+        if config['align_corner'] is None or config['align_corner'] is True:
+            raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")
+        elif config['align_corner'] is False:
+            if config['resolution'] == (8, 1):
+                return s2dnet_8_1_align(config['resnetfpn'])
+    else:
+        raise ValueError(f"LOFTR.BACKBONE_TYPE {config['backbone_type']} not supported.")