vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/XoFTR/configs/data/base.py

@@ -0,0 +1,35 @@
+"""
+The data config will be the last one merged into the main config.
+Setups in data configs will override all existed setups!
+"""
+
+from yacs.config import CfgNode as CN
+_CN = CN()
+_CN.DATASET = CN()
+_CN.TRAINER = CN()
+
+# training data config
+_CN.DATASET.TRAIN_DATA_ROOT = None
+_CN.DATASET.TRAIN_POSE_ROOT = None
+_CN.DATASET.TRAIN_NPZ_ROOT = None
+_CN.DATASET.TRAIN_LIST_PATH = None
+_CN.DATASET.TRAIN_INTRINSIC_PATH = None
+# validation set config
+_CN.DATASET.VAL_DATA_ROOT = None
+_CN.DATASET.VAL_POSE_ROOT = None
+_CN.DATASET.VAL_NPZ_ROOT = None
+_CN.DATASET.VAL_LIST_PATH = None
+_CN.DATASET.VAL_INTRINSIC_PATH = None
+
+# testing data config
+_CN.DATASET.TEST_DATA_ROOT = None
+_CN.DATASET.TEST_POSE_ROOT = None
+_CN.DATASET.TEST_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = None
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# dataset config
+_CN.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.4
+_CN.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0  # for both test and val
+
+cfg = _CN

vismatch/third_party/XoFTR/configs/data/megadepth_trainval_840.py

@@ -0,0 +1,22 @@
+from configs.data.base import cfg
+
+
+TRAIN_BASE_PATH = "data/megadepth/index"
+cfg.DATASET.TRAINVAL_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TRAIN_DATA_ROOT = "data/megadepth/train"
+cfg.DATASET.TRAIN_NPZ_ROOT = f"{TRAIN_BASE_PATH}/scene_info_0.1_0.7"
+cfg.DATASET.TRAIN_LIST_PATH = f"{TRAIN_BASE_PATH}/trainvaltest_list/train_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.0
+
+TEST_BASE_PATH = "data/megadepth/index"
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.VAL_DATA_ROOT = cfg.DATASET.TEST_DATA_ROOT = "data/megadepth/test"
+cfg.DATASET.VAL_NPZ_ROOT = cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}/scene_info_val_1500"
+cfg.DATASET.VAL_LIST_PATH = cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/trainvaltest_list/val_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0   # for both test and val
+
+# 368 scenes in total for MegaDepth
+# (with difficulty balanced (further split each scene to 3 sub-scenes))
+cfg.TRAINER.N_SAMPLES_PER_SUBSET = 100
+
+cfg.DATASET.MGDPT_IMG_RESIZE = 840  # for training on 32GB meme GPUs

vismatch/third_party/XoFTR/configs/data/megadepth_vistir_trainval_640.py

@@ -0,0 +1,23 @@
+from configs.data.base import cfg
+
+
+TRAIN_BASE_PATH = "data/megadepth/index"
+cfg.DATASET.TRAIN_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TRAIN_DATA_ROOT = "data/megadepth/train"
+cfg.DATASET.TRAIN_NPZ_ROOT = f"{TRAIN_BASE_PATH}/scene_info_0.1_0.7"
+cfg.DATASET.TRAIN_LIST_PATH = f"{TRAIN_BASE_PATH}/trainvaltest_list/train_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.0
+
+VAL_BASE_PATH = "data/METU_VisTIR/index"
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.VAL_DATA_SOURCE = "VisTir"
+cfg.DATASET.VAL_DATA_ROOT = cfg.DATASET.TEST_DATA_ROOT = "data/METU_VisTIR"
+cfg.DATASET.VAL_NPZ_ROOT = cfg.DATASET.TEST_NPZ_ROOT = f"{VAL_BASE_PATH}/scene_info_val"
+cfg.DATASET.VAL_LIST_PATH = cfg.DATASET.TEST_LIST_PATH = f"{VAL_BASE_PATH}/val_test_list/val_list.txt"
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0   # for both test and val
+
+# 368 scenes in total for MegaDepth
+# (with difficulty balanced (further split each scene to 3 sub-scenes))
+cfg.TRAINER.N_SAMPLES_PER_SUBSET = 100
+
+cfg.DATASET.MGDPT_IMG_RESIZE = 640  # for training on 11GB mem GPUs

vismatch/third_party/XoFTR/configs/data/pretrain.py

@@ -0,0 +1,8 @@
+from configs.data.base import cfg
+
+cfg.DATASET.TRAIN_DATA_SOURCE = "KAIST"
+cfg.DATASET.TRAIN_DATA_ROOT = "data/kaist-cvpr15"
+cfg.DATASET.VAL_DATA_SOURCE = "KAIST"
+cfg.DATASET.VAL_DATA_ROOT = cfg.DATASET.TEST_DATA_ROOT = "data/kaist-cvpr15"
+
+cfg.DATASET.PRETRAIN_IMG_RESIZE = 640  

vismatch/third_party/XoFTR/configs/xoftr/outdoor/visible_thermal.py

@@ -0,0 +1,17 @@
+from src.config.default import _CN as cfg
+
+cfg.XOFTR.MATCH_COARSE.MATCH_TYPE = 'dual_softmax'
+
+cfg.TRAINER.CANONICAL_LR = 8e-3
+cfg.TRAINER.WARMUP_STEP = 1875  # 3 epochs
+cfg.TRAINER.WARMUP_RATIO = 0.1
+cfg.TRAINER.MSLR_MILESTONES = [8, 12, 16, 20, 24, 30, 36, 42]
+
+# pose estimation
+cfg.TRAINER.RANSAC_PIXEL_THR = 1.5
+
+cfg.TRAINER.OPTIMIZER = "adamw"
+cfg.TRAINER.ADAMW_DECAY = 0.1
+cfg.XOFTR.MATCH_COARSE.TRAIN_COARSE_PERCENT = 0.3
+
+cfg.TRAINER.USE_WANDB = True # use weight and biases

vismatch/third_party/XoFTR/configs/xoftr/pretrain/pretrain.py

@@ -0,0 +1,12 @@
+from src.config.default import _CN as cfg
+
+cfg.TRAINER.CANONICAL_LR = 4e-3
+cfg.TRAINER.WARMUP_STEP = 1250  # 2 epochs
+cfg.TRAINER.WARMUP_RATIO = 0.1
+cfg.TRAINER.MSLR_MILESTONES = [4, 6, 8, 10, 12, 14, 16, 18]
+
+cfg.TRAINER.OPTIMIZER = "adamw"
+cfg.TRAINER.ADAMW_DECAY = 0.1
+
+cfg.TRAINER.USE_WANDB = True # use weight and biases
+

vismatch/third_party/XoFTR/pretrain.py

@@ -0,0 +1,125 @@
+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_pretrain import PretrainDataModule
+from src.lightning.lightning_xoftr_pretrain import PL_XoFTR_Pretrain
+
+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')
+    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_Pretrain(config, pretrained_ckpt=args.ckpt_path, profiler=profiler)
+    loguru_logger.info(f"XoFTR LightningModule initialized!")
+    
+    # lightning data
+    data_module = PretrainDataModule(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(verbose=True, save_top_k=-1,
+                                    save_last=True,
+                                    dirpath=str(ckpt_dir),
+                                    filename='{epoch}')
+    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/XoFTR/src/config/default.py

@@ -0,0 +1,203 @@
+from yacs.config import CfgNode as CN
+
+INFERENCE = False
+
+_CN = CN()
+
+##############  ↓  XoFTR Pipeline  ↓  ##############
+_CN.XOFTR = CN()
+_CN.XOFTR.RESOLUTION = (8, 2)  # options: [(8, 2)]
+_CN.XOFTR.FINE_WINDOW_SIZE = 5  # window_size in fine_level, must be odd
+_CN.XOFTR.MEDIUM_WINDOW_SIZE = 3  # window_size in fine_level, must be odd
+
+# 1. XoFTR-backbone (local feature CNN) config
+_CN.XOFTR.RESNET = CN()
+_CN.XOFTR.RESNET.INITIAL_DIM = 128
+_CN.XOFTR.RESNET.BLOCK_DIMS = [128, 196, 256]  # s1, s2, s3
+
+# 2. XoFTR-coarse module config
+_CN.XOFTR.COARSE = CN()
+_CN.XOFTR.COARSE.INFERENCE = INFERENCE
+_CN.XOFTR.COARSE.D_MODEL = 256
+_CN.XOFTR.COARSE.D_FFN = 256
+_CN.XOFTR.COARSE.NHEAD = 8
+_CN.XOFTR.COARSE.LAYER_NAMES = ['self', 'cross'] * 4
+_CN.XOFTR.COARSE.ATTENTION = 'linear'  # options: ['linear', 'full']
+
+# 3. Coarse-Matching config
+_CN.XOFTR.MATCH_COARSE = CN()
+_CN.XOFTR.MATCH_COARSE.INFERENCE = INFERENCE
+_CN.XOFTR.MATCH_COARSE.D_MODEL = 256
+_CN.XOFTR.MATCH_COARSE.THR = 0.3
+_CN.XOFTR.MATCH_COARSE.BORDER_RM = 2
+_CN.XOFTR.MATCH_COARSE.MATCH_TYPE = 'dual_softmax'  # options: ['dual_softmax']
+_CN.XOFTR.MATCH_COARSE.DSMAX_TEMPERATURE = 0.1
+_CN.XOFTR.MATCH_COARSE.TRAIN_COARSE_PERCENT = 0.2  # training tricks: save GPU memory
+_CN.XOFTR.MATCH_COARSE.TRAIN_PAD_NUM_GT_MIN = 200  # training tricks: avoid DDP deadlock
+
+# 4. XoFTR-fine module config
+_CN.XOFTR.FINE = CN()
+_CN.XOFTR.FINE.DENSER = False # if true, match all features in fine-level windows
+_CN.XOFTR.FINE.INFERENCE = INFERENCE
+_CN.XOFTR.FINE.DSMAX_TEMPERATURE = 0.1
+_CN.XOFTR.FINE.THR = 0.1
+_CN.XOFTR.FINE.MLP_HIDDEN_DIM_COEF = 2 # coef for  mlp hidden dim (hidden_dim = feat_dim * coef)
+_CN.XOFTR.FINE.NHEAD_FINE_LEVEL = 8
+_CN.XOFTR.FINE.NHEAD_MEDIUM_LEVEL = 7
+
+
+# 5. XoFTR Losses
+
+_CN.XOFTR.LOSS = CN()
+_CN.XOFTR.LOSS.FOCAL_ALPHA = 0.25
+_CN.XOFTR.LOSS.FOCAL_GAMMA = 2.0
+_CN.XOFTR.LOSS.POS_WEIGHT = 1.0
+_CN.XOFTR.LOSS.NEG_WEIGHT = 1.0
+
+# -- # coarse-level
+_CN.XOFTR.LOSS.COARSE_WEIGHT = 0.5
+# -- # fine-level
+_CN.XOFTR.LOSS.FINE_WEIGHT = 0.3
+# -- # sub-pixel
+_CN.XOFTR.LOSS.SUB_WEIGHT = 1 * 10**4
+
+##############  Dataset  ##############
+_CN.DATASET = CN()
+# 1. data config
+# training and validating
+_CN.DATASET.TRAIN_DATA_SOURCE = None  # options: ['ScanNet', 'MegaDepth']
+_CN.DATASET.TRAIN_DATA_ROOT = None
+_CN.DATASET.TRAIN_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.TRAIN_NPZ_ROOT = None
+_CN.DATASET.TRAIN_LIST_PATH = None
+_CN.DATASET.TRAIN_INTRINSIC_PATH = None
+_CN.DATASET.VAL_DATA_SOURCE = None
+_CN.DATASET.VAL_DATA_ROOT = None
+_CN.DATASET.VAL_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.VAL_NPZ_ROOT = None
+_CN.DATASET.VAL_LIST_PATH = None    # None if val data from all scenes are bundled into a single npz file
+_CN.DATASET.VAL_INTRINSIC_PATH = None
+# testing
+_CN.DATASET.TEST_DATA_SOURCE = None
+_CN.DATASET.TEST_DATA_ROOT = None
+_CN.DATASET.TEST_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.TEST_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = None   # None if test data from all scenes are bundled into a single npz file
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# 2. dataset config
+# general options
+_CN.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.4  # discard data with overlap_score < min_overlap_score
+_CN.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0
+_CN.DATASET.AUGMENTATION_TYPE = "rgb_thermal"  # options: [None, 'dark', 'mobile']
+
+# MegaDepth options
+_CN.DATASET.MGDPT_IMG_RESIZE = 640  # resize the longer side, zero-pad bottom-right to square.
+_CN.DATASET.MGDPT_IMG_PAD = True  # pad img to square with size = MGDPT_IMG_RESIZE
+_CN.DATASET.MGDPT_DEPTH_PAD = True  # pad depthmap to square with size = 2000
+_CN.DATASET.MGDPT_DF = 8
+
+# VisTir options
+_CN.DATASET.VISTIR_IMG_RESIZE = 640  # resize the longer side, zero-pad bottom-right to square.
+_CN.DATASET.VISTIR_IMG_PAD = False  # pad img to square with size = VISTIR_IMG_RESIZE
+_CN.DATASET.VISTIR_DF = 8
+
+# Pretrain dataset options
+_CN.DATASET.PRETRAIN_IMG_RESIZE = 640  # resize the longer side, zero-pad bottom-right to square.
+_CN.DATASET.PRETRAIN_IMG_PAD = True  # pad img to square with size = PRETRAIN_IMG_RESIZE
+_CN.DATASET.PRETRAIN_DF = 8
+_CN.DATASET.PRETRAIN_FRAME_GAP = 2 # the gap between video frames of Kaist dataset
+
+##############  Trainer  ##############
+_CN.TRAINER = CN()
+_CN.TRAINER.WORLD_SIZE = 1
+_CN.TRAINER.CANONICAL_BS = 64
+_CN.TRAINER.CANONICAL_LR = 6e-3
+_CN.TRAINER.SCALING = None  # this will be calculated automatically
+_CN.TRAINER.FIND_LR = False  # use learning rate finder from pytorch-lightning
+
+_CN.TRAINER.USE_WANDB = False # use weight and biases
+
+# optimizer
+_CN.TRAINER.OPTIMIZER = "adamw"  # [adam, adamw]
+_CN.TRAINER.TRUE_LR = None  # this will be calculated automatically at runtime
+_CN.TRAINER.ADAM_DECAY = 0.  # ADAM: for adam
+_CN.TRAINER.ADAMW_DECAY = 0.1
+
+# step-based warm-up
+_CN.TRAINER.WARMUP_TYPE = 'linear'  # [linear, constant]
+_CN.TRAINER.WARMUP_RATIO = 0.
+_CN.TRAINER.WARMUP_STEP = 4800
+
+# learning rate scheduler
+_CN.TRAINER.SCHEDULER = 'MultiStepLR'  # [MultiStepLR, CosineAnnealing, ExponentialLR]
+_CN.TRAINER.SCHEDULER_INTERVAL = 'epoch'    # [epoch, step]
+_CN.TRAINER.MSLR_MILESTONES = [3, 6, 9, 12]  # MSLR: MultiStepLR
+_CN.TRAINER.MSLR_GAMMA = 0.5
+_CN.TRAINER.COSA_TMAX = 30  # COSA: CosineAnnealing
+_CN.TRAINER.ELR_GAMMA = 0.999992  # ELR: ExponentialLR, this value for 'step' interval
+
+# plotting related
+_CN.TRAINER.ENABLE_PLOTTING = True
+_CN.TRAINER.N_VAL_PAIRS_TO_PLOT = 128     # number of val/test paris for plotting
+_CN.TRAINER.PLOT_MODE = 'evaluation'  # ['evaluation', 'confidence']
+_CN.TRAINER.PLOT_MATCHES_ALPHA = 'dynamic'
+
+# geometric metrics and pose solver
+_CN.TRAINER.EPI_ERR_THR = 5e-4  # recommendation: 5e-4 for ScanNet, 1e-4 for MegaDepth (from SuperGlue)
+_CN.TRAINER.POSE_GEO_MODEL = 'E'  # ['E', 'F', 'H']
+_CN.TRAINER.POSE_ESTIMATION_METHOD = 'RANSAC'  # [RANSAC, DEGENSAC, MAGSAC]
+_CN.TRAINER.RANSAC_PIXEL_THR = 0.5
+_CN.TRAINER.RANSAC_CONF = 0.99999
+_CN.TRAINER.RANSAC_MAX_ITERS = 10000
+_CN.TRAINER.USE_MAGSACPP = False
+
+# data sampler for train_dataloader
+_CN.TRAINER.DATA_SAMPLER = 'scene_balance'  # options: ['scene_balance', 'random', 'normal']
+# 'scene_balance' config
+_CN.TRAINER.N_SAMPLES_PER_SUBSET = 200
+_CN.TRAINER.SB_SUBSET_SAMPLE_REPLACEMENT = True  # whether sample each scene with replacement or not
+_CN.TRAINER.SB_SUBSET_SHUFFLE = True  # after sampling from scenes, whether shuffle within the epoch or not
+_CN.TRAINER.SB_REPEAT = 1  # repeat N times for training the sampled data
+# 'random' config
+_CN.TRAINER.RDM_REPLACEMENT = True
+_CN.TRAINER.RDM_NUM_SAMPLES = None
+
+# gradient clipping
+_CN.TRAINER.GRADIENT_CLIPPING = 0.5
+
+# reproducibility
+# This seed affects the data sampling. With the same seed, the data sampling is promised
+# to be the same. When resume training from a checkpoint, it's better to use a different
+# seed, otherwise the sampled data will be exactly the same as before resuming, which will
+# cause less unique data items sampled during the entire training.
+# Use of different seed values might affect the final training result, since not all data items
+# are used during training on ScanNet. (60M pairs of images sampled during traing from 230M pairs in total.)
+_CN.TRAINER.SEED = 66
+
+##############  Pretrain  ##############
+_CN.PRETRAIN = CN()
+_CN.PRETRAIN.PATCH_SIZE = 64 # patch sıze for masks
+_CN.PRETRAIN.MASK_RATIO = 0.5 
+_CN.PRETRAIN.MAE_MARGINS = [0, 0.4, 0, 0] # margins not to be masked (up bottom left right)
+_CN.PRETRAIN.VAL_SEED = 42 # rng seed to crate the same masks for validation
+
+_CN.XOFTR.PRETRAIN_PATCH_SIZE = _CN.PRETRAIN.PATCH_SIZE 
+
+##############  Test/Inference  ##############
+_CN.TEST = CN()
+_CN.TEST.IMG0_RESIZE = 640 # resize the longer side
+_CN.TEST.IMG1_RESIZE = 640 # resize the longer side
+_CN.TEST.DF = 8
+_CN.TEST.PADDING = False  # pad img to square with size = IMG0_RESIZE, IMG1_RESIZE
+_CN.TEST.COARSE_SCALE = 0.125
+
+def get_cfg_defaults(inference=False):
+    """Get a yacs CfgNode object with default values for my_project."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    if inference:
+        _CN.XOFTR.COARSE.INFERENCE = True
+        _CN.XOFTR.MATCH_COARSE.INFERENCE = True
+        _CN.XOFTR.FINE.INFERENCE = True
+    return _CN.clone()

vismatch/third_party/XoFTR/src/datasets/megadepth.py

@@ -0,0 +1,143 @@
+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 src.utils.dataset import read_megadepth_gray, read_megadepth_depth
+
+def correct_image_paths(scene_info):
+    """Changes the path format from undistorted images from D2Net to MegaDepth_v1 format"""
+    image_paths = scene_info["image_paths"]
+    for ii in range(len(image_paths)):
+        if image_paths[ii] is not None:
+            folds = image_paths[ii].split("/")
+            path = osp.join("phoenix/S6/zl548/MegaDepth_v1/", folds[1], "dense0/imgs", folds[3] )
+            image_paths[ii] = path
+    scene_info["image_paths"] = image_paths
+    return scene_info
+
+class MegaDepthDataset(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,
+                 **kwargs):
+        """
+        Manage one scene(npz_path) of MegaDepth dataset.
+        
+        Args:
+            root_dir (str): megadepth root directory that has `phoenix`.
+            npz_path (str): {scene_id}.npz path. This contains image pair information of a scene.
+            mode (str): options are ['train', 'val', 'test']
+            min_overlap_score (float): how much a pair should have in common. In range of [0, 1]. Set to 0 when testing.
+            img_resize (int, optional): the longer edge of resized images. None for no resize. 640 is recommended.
+                                        This is useful during training with batches and testing with memory intensive algorithms.
+            df (int, optional): image size division factor. NOTE: this will change the final image size after img_resize.
+            img_padding (bool): If set to 'True', zero-pad the image to squared size. This is useful during training.
+            depth_padding (bool): If set to 'True', zero-pad depthmap to (2000, 2000). This is useful during training.
+            augment_fn (callable, optional): augments images with pre-defined visual effects.
+        """
+        super().__init__()
+        self.root_dir = root_dir
+        self.mode = mode
+        self.scene_id = npz_path.split('.')[0]
+
+        # 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 = 0
+        self.scene_info = np.load(npz_path, allow_pickle=True)
+        self.scene_info = correct_image_paths(self.scene_info)
+        self.pair_infos = self.scene_info['pair_infos'].copy()
+        del self.scene_info['pair_infos']
+        self.pair_infos = [pair_info for pair_info in self.pair_infos if pair_info[1] > min_overlap_score]
+
+        # parameters for image resizing, padding and depthmap padding
+        if mode == 'train':
+            assert img_resize is not None and img_padding and depth_padding
+        self.img_resize = img_resize
+        self.df = df
+        self.img_padding = img_padding
+        self.depth_max_size = 2000 if depth_padding else None  # the upperbound of depthmaps size in megadepth.
+
+        # for training XoFTR
+        # self.augment_fn = augment_fn if mode == 'train' else None
+        self.augment_fn = augment_fn
+        self.coarse_scale = getattr(kwargs, 'coarse_scale', 0.125)
+
+    def __len__(self):
+        return len(self.pair_infos)
+
+    def __getitem__(self, idx):
+        (idx0, idx1), overlap_score, central_matches = self.pair_infos[idx]
+
+        # read grayscale image and mask. (1, h, w) and (h, w)
+        img_name0 = osp.join(self.root_dir, self.scene_info['image_paths'][idx0])
+        img_name1 = osp.join(self.root_dir, self.scene_info['image_paths'][idx1])
+        
+        if getattr(self.augment_fn, 'random_switch', False):
+            im_num = torch.randint(0, 2, (1,))
+            augment_fn_0 = lambda x: self.augment_fn(x, image_num=im_num)
+            augment_fn_1 = lambda x: self.augment_fn(x, image_num=1-im_num)
+        else:
+            augment_fn_0 = self.augment_fn
+            augment_fn_1 = self.augment_fn
+        image0, mask0, scale0 = read_megadepth_gray(
+            img_name0, self.img_resize, self.df, self.img_padding, augment_fn=augment_fn_0)
+        image1, mask1, scale1 = read_megadepth_gray(
+            img_name1, self.img_resize, self.df, self.img_padding, augment_fn=augment_fn_1)
+
+        # read depth. shape: (h, w)
+        if self.mode in ['train', 'val']:
+            depth0 = read_megadepth_depth(
+                osp.join(self.root_dir, self.scene_info['depth_paths'][idx0]), pad_to=self.depth_max_size)
+            depth1 = read_megadepth_depth(
+                osp.join(self.root_dir, self.scene_info['depth_paths'][idx1]), pad_to=self.depth_max_size)
+        else:
+            depth0 = depth1 = torch.tensor([])
+
+        # read intrinsics of original size
+        K_0 = torch.tensor(self.scene_info['intrinsics'][idx0].copy(), dtype=torch.float).reshape(3, 3)
+        K_1 = torch.tensor(self.scene_info['intrinsics'][idx1].copy(), dtype=torch.float).reshape(3, 3)
+
+        # read and compute relative poses
+        T0 = self.scene_info['poses'][idx0]
+        T1 = self.scene_info['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()
+
+        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,
+            'scale0': scale0,  # [scale_w, scale_h]
+            'scale1': scale1,
+            'dataset_name': 'MegaDepth',
+            'scene_id': self.scene_id,
+            'pair_id': idx,
+            'pair_names': (self.scene_info['image_paths'][idx0], self.scene_info['image_paths'][idx1]),
+        }
+
+        # for XoFTR training
+        if mask0 is not None:  # img_padding is True
+            if self.coarse_scale:
+                [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()
+            data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1})
+
+        return data