vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/XoFTR/train.py

@@ -0,0 +1,126 @@
+import math
+import argparse
+import pprint
+from distutils.util import strtobool
+from pathlib import Path
+from loguru import logger as loguru_logger
+from datetime import datetime
+
+import pytorch_lightning as pl
+from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.plugins import DDPPlugin
+
+from src.config.default import get_cfg_defaults
+from src.utils.misc import get_rank_zero_only_logger, setup_gpus
+from src.utils.profiler import build_profiler
+from src.lightning.data import MultiSceneDataModule
+from src.lightning.lightning_xoftr import PL_XoFTR
+
+loguru_logger = get_rank_zero_only_logger(loguru_logger)
+
+
+def parse_args():
+    # init a costum parser which will be added into pl.Trainer parser
+    # check documentation: https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#trainer-flags
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        'data_cfg_path', type=str, help='data config path')
+    parser.add_argument(
+        'main_cfg_path', type=str, help='main config path')
+    parser.add_argument(
+        '--exp_name', type=str, default='default_exp_name')
+    parser.add_argument(
+        '--batch_size', type=int, default=4, help='batch_size per gpu')
+    parser.add_argument(
+        '--num_workers', type=int, default=4)
+    parser.add_argument(
+        '--pin_memory', type=lambda x: bool(strtobool(x)),
+        nargs='?', default=True, help='whether loading data to pinned memory or not')
+    parser.add_argument(
+        '--ckpt_path', type=str, default=None,
+        help='pretrained checkpoint path, helpful for using a pre-trained coarse-only XoFTR')
+    parser.add_argument(
+        '--disable_ckpt', action='store_true',
+        help='disable checkpoint saving (useful for debugging).')
+    parser.add_argument(
+        '--profiler_name', type=str, default=None,
+        help='options: [inference, pytorch], or leave it unset')
+    parser.add_argument(
+        '--parallel_load_data', action='store_true',
+        help='load datasets in with multiple processes.')
+
+    parser = pl.Trainer.add_argparse_args(parser)
+    return parser.parse_args()
+
+
+def main():
+    # parse arguments
+    args = parse_args()
+    rank_zero_only(pprint.pprint)(vars(args))
+
+    # init default-cfg and merge it with the main- and data-cfg
+    config = get_cfg_defaults()
+    config.merge_from_file(args.main_cfg_path)
+    config.merge_from_file(args.data_cfg_path)
+    pl.seed_everything(config.TRAINER.SEED)  # reproducibility
+    
+    # scale lr and warmup-step automatically
+    args.gpus = _n_gpus = setup_gpus(args.gpus)
+    config.TRAINER.WORLD_SIZE = _n_gpus * args.num_nodes
+    config.TRAINER.TRUE_BATCH_SIZE = config.TRAINER.WORLD_SIZE * args.batch_size
+    _scaling = config.TRAINER.TRUE_BATCH_SIZE / config.TRAINER.CANONICAL_BS
+    config.TRAINER.SCALING = _scaling
+    config.TRAINER.TRUE_LR = config.TRAINER.CANONICAL_LR * _scaling
+    config.TRAINER.WARMUP_STEP = math.floor(config.TRAINER.WARMUP_STEP / _scaling)
+    
+    
+    # lightning module
+    profiler = build_profiler(args.profiler_name)
+    model = PL_XoFTR(config, pretrained_ckpt=args.ckpt_path, profiler=profiler)
+    loguru_logger.info(f"XoFTR LightningModule initialized!")
+    
+    # lightning data
+    data_module = MultiSceneDataModule(args, config)
+    loguru_logger.info(f"XoFTR DataModule initialized!")
+    
+    # TensorBoard Logger
+    logger = [TensorBoardLogger(save_dir='logs/tb_logs', name=args.exp_name, default_hp_metric=False)]
+    ckpt_dir = Path(logger[0].log_dir) / 'checkpoints'
+    if config.TRAINER.USE_WANDB:
+        logger.append(WandbLogger(name=args.exp_name + f"_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}",
+                                project='XoFTR'))
+    
+    # Callbacks
+    # TODO: update ModelCheckpoint to monitor multiple metrics
+    ckpt_callback = ModelCheckpoint(monitor='auc@10', verbose=True, save_top_k=-1, mode='max',
+                                    save_last=True,
+                                    dirpath=str(ckpt_dir),
+                                    filename='{epoch}-{auc@5:.3f}-{auc@10:.3f}-{auc@20:.3f}')
+    lr_monitor = LearningRateMonitor(logging_interval='step')
+    callbacks = [lr_monitor]
+    if not args.disable_ckpt:
+        callbacks.append(ckpt_callback)
+    
+    # Lightning Trainer
+    trainer = pl.Trainer.from_argparse_args(
+        args,
+        plugins=DDPPlugin(find_unused_parameters=True,
+                          num_nodes=args.num_nodes,
+                          sync_batchnorm=config.TRAINER.WORLD_SIZE > 0),
+        gradient_clip_val=config.TRAINER.GRADIENT_CLIPPING,
+        callbacks=callbacks,
+        logger=logger,
+        sync_batchnorm=config.TRAINER.WORLD_SIZE > 0,
+        replace_sampler_ddp=False,  # use custom sampler
+        reload_dataloaders_every_epoch=False,  # avoid repeated samples!
+        weights_summary='full',
+        profiler=profiler)
+    loguru_logger.info(f"Trainer initialized!")
+    loguru_logger.info(f"Start training!")
+    trainer.fit(model, datamodule=data_module)
+
+
+if __name__ == '__main__':
+    main()

vismatch/third_party/accelerated_features/hubconf.py

@@ -0,0 +1,15 @@
+dependencies = ['torch']
+from modules.xfeat import XFeat as _XFeat
+import torch
+
+def XFeat(pretrained=True, top_k=4096, detection_threshold=0.05):
+    """
+    XFeat model
+    pretrained (bool): kwargs, load pretrained weights into the model
+    """
+    weights = None
+    if pretrained:
+        weights = torch.hub.load_state_dict_from_url("https://github.com/verlab/accelerated_features/raw/main/weights/xfeat.pt")
+    
+    model = _XFeat(weights, top_k=top_k, detection_threshold=detection_threshold)
+    return model

vismatch/third_party/accelerated_features/minimal_example.py

@@ -0,0 +1,49 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+    Minimal example of how to use XFeat.
+"""
+
+import numpy as np
+import os
+import torch
+import tqdm
+
+from modules.xfeat import XFeat
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '' #Force CPU, comment for GPU
+
+xfeat = XFeat()
+
+#Random input
+x = torch.randn(1,3,480,640)
+
+#Simple inference with batch = 1
+output = xfeat.detectAndCompute(x, top_k = 4096)[0]
+print("----------------")
+print("keypoints: ", output['keypoints'].shape)
+print("descriptors: ", output['descriptors'].shape)
+print("scores: ", output['scores'].shape)
+print("----------------\n")
+
+x = torch.randn(1,3,480,640)
+# Stress test
+for i in tqdm.tqdm(range(100), desc="Stress test on VGA resolution"):
+	output = xfeat.detectAndCompute(x, top_k = 4096)
+
+# Batched mode
+x = torch.randn(4,3,480,640)
+outputs = xfeat.detectAndCompute(x, top_k = 4096)
+print("# detected features on each batch item:", [len(o['keypoints']) for o in outputs])
+
+# Match two images with sparse features
+x1 = torch.randn(1,3,480,640)
+x2 = torch.randn(1,3,480,640)
+mkpts_0, mkpts_1 = xfeat.match_xfeat(x1, x2)
+
+# Match two images with semi-dense approach -- batched mode with batch size 4
+x1 = torch.randn(4,3,480,640)
+x2 = torch.randn(4,3,480,640)
+matches_list = xfeat.match_xfeat_star(x1, x2)
+print(matches_list[0].shape)

vismatch/third_party/accelerated_features/modules/__init__.py

@@ -0,0 +1,4 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+"""

vismatch/third_party/accelerated_features/modules/dataset/__init__.py

@@ -0,0 +1,5 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+"""

vismatch/third_party/accelerated_features/modules/dataset/augmentation.py

@@ -0,0 +1,314 @@
+"""
+    "XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+    https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+    This script implements color + geometric transformations using Kornia.
+    Given a dataset of random real unlabeled images, we apply photometric transformations, 
+    homography warps and also TPS warps. It also handles black borders by
+    pasting a random background image.
+"""
+
+import torch
+from torch import nn
+from torch.utils.data import Dataset
+import torch.utils.data as data
+from torchvision import transforms
+import torch.nn.functional as F
+
+import cv2
+import kornia
+import kornia.augmentation as K
+from kornia.geometry.transform import get_tps_transform as findTPS
+from kornia.geometry.transform import warp_points_tps, warp_image_tps
+
+import glob
+import random
+import tqdm
+
+import numpy as np
+import pdb
+import time
+
+random.seed(0)
+torch.manual_seed(0)
+
+def generateRandomTPS(shape, grid = (8, 6), GLOBAL_MULTIPLIER = 0.3, prob = 0.5):
+
+    h, w = shape
+    sh, sw = h/grid[0], w/grid[1]
+    src = torch.dstack(torch.meshgrid(torch.arange(0, h + sh , sh),
+                         torch.arange(0, w + sw , sw), indexing='ij'))
+
+    offsets = torch.rand(grid[0]+1, grid[1]+1, 2) - 0.5
+    offsets *= torch.tensor([ sh/2, sw/2 ]).view(1, 1, 2)  * min(0.97, 2.0 * GLOBAL_MULTIPLIER)
+    dst = src + offsets if np.random.uniform() < prob else src
+    
+    src, dst = src.view(1, -1, 2), dst.view(1, -1, 2)
+    src = (src / torch.tensor([h,w]).view(1,1,2) ) * 2 - 1.
+    dst = (dst / torch.tensor([h,w]).view(1,1,2) ) * 2 - 1.
+    weights, A = findTPS(dst, src)
+
+    return src, weights, A
+
+
+def generateRandomHomography(shape, GLOBAL_MULTIPLIER = 0.3):
+    #Generate random in-plane rotation [-theta,+theta]
+    theta = np.radians(np.random.uniform(-30, 30))
+
+    #Generate random scale in both x and y
+    scale_x, scale_y = np.random.uniform(0.35, 1.2, 2)
+
+    #Generate random translation shift
+    tx , ty = -shape[1]/2.0 , -shape[0]/2.0 
+    txn, tyn = np.random.normal(0, 120.0*GLOBAL_MULTIPLIER, 2) 
+
+    c, s = np.cos(theta), np.sin(theta)
+
+    #Affine coeffs
+    sx , sy = np.random.normal(0,0.6*GLOBAL_MULTIPLIER,2)
+
+    #Projective coeffs
+    p1 , p2 = np.random.normal(0,0.006*GLOBAL_MULTIPLIER,2)
+
+
+    # Build Homography from parmeterizations
+    H_t = np.array(((1,0, tx), (0, 1, ty), (0,0,1))) #t
+
+    H_r = np.array(((c,-s, 0), (s, c, 0), (0,0,1))) #rotation,
+    H_a = np.array(((1,sy, 0), (sx, 1, 0), (0,0,1))) # affine
+    H_p = np.array(((1, 0, 0), (0 , 1, 0), (p1,p2,1))) # projective
+
+    H_s = np.array(((scale_x,0, 0), (0, scale_y, 0), (0,0,1))) #scale
+    H_b = np.array(((1.0,0,-tx +txn), (0, 1, -ty + tyn), (0,0,1))) #t_back,
+
+    #H = H_e * H_s * H_a * H_p
+    H = np.dot(np.dot(np.dot(np.dot(np.dot(H_b,H_s),H_p),H_a),H_r),H_t)
+
+    return H
+
+
+class AugmentationPipe(nn.Module):
+    def __init__(
+                self, device, load_dataset = True,
+                img_dir = "/homeLocal/guipotje/sshfs/datasets/coco_20k",
+                warp_resolution = (1200, 900),
+                out_resolution = (400, 300),
+                sides_crop = 0.2,
+                max_num_imgs = 50,
+                num_test_imgs = 10,
+                batch_size = 1,
+                photometric = True,
+                geometric = True,
+                reload_step = 1_000
+                ):
+        super(AugmentationPipe, self).__init__()
+        self.half = 16
+        self.device = device
+
+        self.dims = warp_resolution
+        self.batch_size = batch_size
+        self.out_resolution = out_resolution 
+        self.sides_crop = sides_crop
+        self.max_num_imgs = max_num_imgs
+        self.num_test_imgs = num_test_imgs
+        self.dims_t = torch.tensor([int(self.dims[0]*(1. - self.sides_crop)) - int(self.dims[0]*self.sides_crop) -1,
+                                    int(self.dims[1]*(1. - self.sides_crop)) - int(self.dims[1]*self.sides_crop) -1]).float().to(device).view(1,1,2)
+        self.dims_s = torch.tensor([ self.dims_t[0,0,0] / out_resolution[0],
+                                     self.dims_t[0,0,1] / out_resolution[1]]).float().to(device).view(1,1,2) 
+
+        self.all_imgs = glob.glob(img_dir + '/*.jpg') + glob.glob(img_dir + '/*.png')
+        
+        self.photometric = photometric
+        self.geometric = geometric
+        self.cnt = 1
+        self.reload_step = reload_step
+
+        list_augmentation = [
+                        #kornia.augmentation.RandomChannelShuffle(p=0.5),
+                        kornia.augmentation.ColorJitter(0.15, 0.15, 0.15, 0.15, p=1.),
+                        kornia.augmentation.RandomEqualize(p = 0.4),
+                        kornia.augmentation.RandomGaussianBlur(p = 0.3, sigma = (2.0, 2.0), kernel_size = (7,7))
+                        ]
+        if photometric is False:
+            list_augmentation = []
+
+        self.aug_list = kornia.augmentation.ImageSequential(*list_augmentation)
+        
+        if len(self.all_imgs) < 10:
+            raise RuntimeError('Couldnt find enough images to train. Please check the path: ', img_dir)
+
+        if load_dataset:
+            print('[Synthetic] Found a total of ', len(self.all_imgs), ' images for training..')
+
+            if len(self.all_imgs) - num_test_imgs < max_num_imgs:
+                raise RuntimeError('Error: test set overlaps with training set! Decrease number of test imgs')
+
+            self.load_imgs()
+
+            self.TPS = True
+
+
+    def load_imgs(self):
+
+            random.shuffle(self.all_imgs)
+
+            train = []
+            fast = cv2.FastFeatureDetector_create(30)
+            for p in tqdm.tqdm(self.all_imgs[:self.max_num_imgs], desc='loading train'):
+                im = cv2.imread(p)
+                halfH, halfW = im.shape[0]//2, im.shape[1]//2
+                if halfH > halfW:
+                    im = np.rot90(im)
+                    halfH, halfW = halfW, halfH
+
+                if im.shape[0] != self.dims[1] or im.shape[1] != self.dims[0]:
+                    im = cv2.resize(im, self.dims)
+
+                train.append(np.copy(im))
+
+            self.train = train
+            
+            self.test = [
+                        cv2.resize(cv2.imread(p), self.dims)                 
+                        for p in tqdm.tqdm(self.all_imgs[-self.num_test_imgs:],
+                                           desc='loading test')
+                        ] 
+
+    def norm_pts_grid(self, x):
+        if len(x.size()) == 2:
+            return (x.view(1,-1,2) * self.dims_s / self.dims_t) * 2. - 1 
+        return (x * self.dims_s / self.dims_t) * 2. - 1
+
+    def denorm_pts_grid(self, x):
+        if len(x.size()) == 2:
+            return ((x.view(1,-1,2) + 1) / 2.) / self.dims_s * self.dims_t
+        return ((x+1) / 2.) / self.dims_s * self.dims_t
+
+    def rnd_kps(self, shape, n = 256):
+        h, w = shape
+        kps = torch.rand(size = (3,n)).to(self.device)
+        kps[0,:]*=w
+        kps[1,:]*=h
+        kps[2,:] = 1.0
+
+        return kps
+
+    def warp_points(self, H, pts):
+      scale = self.dims_s.view(-1,2)
+      offset = torch.tensor([int(self.dims[0]*self.sides_crop), int(self.dims[1]*self.sides_crop)], device = pts.device).float()
+      pts = pts*scale + offset
+      pts = torch.vstack( [pts.t(), torch.ones(1, pts.shape[0], device = pts.device)])
+      warped = torch.matmul(H, pts)
+      warped = warped / warped[2,...] 
+      warped = warped.t()[:, :2]
+      return (warped - offset) / scale
+
+    @torch.inference_mode()
+    def forward(self, x, difficulty = 0.3, TPS = False, prob_deformation = 0.5, test = False):
+        """
+            Perform augmentation to a batch of images.
+
+            input:
+                x -> torch.Tensor(B, C, H, W): rgb images
+                difficulty -> float: level of difficulty, 0.1 is medium, 0.3 is already pretty hard
+                tps -> bool: Wether to apply non-rigid deformations in images
+                prob_deformation -> float: probability to apply a deformation
+
+            return:
+                'output'    ->   torch.Tensor(B, C, H, W): rgb images
+                Tuple:
+                    'H'       ->   torch.Tensor(3,3): homography matrix 
+                    'mask'  ->     torch.Tensor(B, H, W): mask of valid pixels after warp
+                    (deformation only)
+                    src, weights, A are parameters from a TPS warp (all torch.Tensors)
+
+        """
+
+        if self.cnt % self.reload_step == 0:
+            self.load_imgs()
+
+        if self.geometric is False:
+            difficulty = 0.
+
+        with torch.no_grad():
+            x = (x/255.).to(self.device)
+            b, c, h, w = x.shape
+            shape = (h, w)
+              
+            ######## Geometric Transformations
+
+            H = torch.tensor(np.array([generateRandomHomography(shape, difficulty) for b in range(self.batch_size)]),
+                               dtype = torch.float32).to(self.device)
+            
+            output = kornia.geometry.transform.warp_perspective(x, H,
+                            dsize = shape, padding_mode = 'zeros')
+
+            #crop % of image boundaries each side to reduce invalid pixels after warps
+            low_h = int(h * self.sides_crop); low_w = int(w*self.sides_crop)
+            high_h = int(h*(1. - self.sides_crop)); high_w= int(w * (1. - self.sides_crop))
+            output = output[..., low_h:high_h, low_w:high_w]
+            x = x[..., low_h:high_h, low_w:high_w]
+
+            #apply TPS if desired:
+            if TPS:
+                src, weights, A = None, None, None
+                for b in range(self.batch_size):
+                    b_src, b_weights, b_A = generateRandomTPS(shape, (8,6), difficulty, prob = prob_deformation)
+                    b_src, b_weights, b_A = b_src.to(self.device), b_weights.to(self.device), b_A.to(self.device)
+
+                    if src is None:
+                        src, weights, A = b_src, b_weights, b_A
+                    else:
+                        src = torch.cat((b_src, src))
+                        weights = torch.cat((b_weights, weights))
+                        A = torch.cat((b_A, A))
+
+                output = warp_image_tps(output, src, weights, A)
+
+            output = F.interpolate(output, self.out_resolution[::-1], mode = 'nearest')
+            x = F.interpolate(x, self.out_resolution[::-1], mode = 'nearest')
+
+            mask = ~torch.all(output == 0, dim=1, keepdim=True)
+            mask = mask.expand(-1,3,-1,-1)
+
+            # Make-up invalid regions with texture from the batch
+            rv = 1 if not TPS else 2
+            output_shifted = torch.roll(x, rv, 0)
+            output[~mask] = output_shifted[~mask]
+            mask = mask[:, 0, :, :]
+
+            ######## Photometric Transformations
+            output = self.aug_list(output)
+
+            b, c, h, w = output.shape
+            #Correlated Gaussian Noise
+            if np.random.uniform() > 0.5 and self.photometric:
+                noise = F.interpolate(torch.randn_like(output)*(10/255), (h//2, w//2))
+                noise = F.interpolate(noise, (h, w), mode = 'bicubic')
+                output = torch.clip( output + noise, 0., 1.)
+
+            #Random shadows
+            if np.random.uniform() > 0.6 and self.photometric:
+                noise = torch.rand((b, 1, h//64, w//64), device = self.device) * 1.3
+                noise = torch.clip(noise, 0.25, 1.0)
+                noise = F.interpolate(noise, (h, w), mode = 'bicubic')
+                noise = noise.expand(-1, 3, -1, -1)
+                output *= noise
+                output = torch.clip( output, 0., 1.)
+
+            self.cnt+=1
+
+        if TPS:
+            return output, (H, src, weights, A, mask)
+        else:
+            return output, (H, mask)
+
+    def get_correspondences(self, kps_target, T):
+        H, H2, src, W, A = T
+        undeformed  = self.denorm_pts_grid(   
+                                        warp_points_tps(self.norm_pts_grid(kps_target),
+                                        src, W, A) ).view(-1,2)
+
+        warped_to_src = self.warp_points(H@torch.inverse(H2), undeformed)
+
+        return warped_to_src

vismatch/third_party/accelerated_features/modules/dataset/megadepth/__init__.py

@@ -0,0 +1,7 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+    MegaDepth data handling was adapted from 
+    LoFTR official code: https://github.com/zju3dv/LoFTR/blob/master/src/datasets/megadepth.py
+"""