Spaces:
Running
Running
| import pytorch_lightning as pl | |
| import argparse | |
| import pprint | |
| from loguru import logger as loguru_logger | |
| from src.config.default import get_cfg_defaults | |
| from src.utils.profiler import build_profiler | |
| from src.lightning.data import MultiSceneDataModule | |
| from src.lightning.lightning_loftr import PL_LoFTR | |
| import torch | |
| 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( | |
| '--ckpt_path', type=str, default="weights/indoor_ds.ckpt", help='path to the checkpoint') | |
| parser.add_argument( | |
| '--dump_dir', type=str, default=None, help="if set, the matching results will be dump to dump_dir") | |
| parser.add_argument( | |
| '--profiler_name', type=str, default=None, help='options: [inference, pytorch], or leave it unset') | |
| parser.add_argument( | |
| '--batch_size', type=int, default=1, help='batch_size per gpu') | |
| parser.add_argument( | |
| '--num_workers', type=int, default=2) | |
| parser.add_argument( | |
| '--thr', type=float, default=None, help='modify the coarse-level matching threshold.') | |
| parser.add_argument( | |
| '--pixel_thr', type=float, default=None, help='modify the RANSAC threshold.') | |
| parser.add_argument( | |
| '--ransac', type=str, default=None, help='modify the RANSAC method') | |
| parser.add_argument( | |
| '--scannetX', type=int, default=None, help='ScanNet resize X') | |
| parser.add_argument( | |
| '--scannetY', type=int, default=None, help='ScanNet resize Y') | |
| parser.add_argument( | |
| '--megasize', type=int, default=None, help='MegaDepth resize') | |
| parser.add_argument( | |
| '--npe', action='store_true', default=False, help='') | |
| parser.add_argument( | |
| '--fp32', action='store_true', default=False, help='') | |
| parser.add_argument( | |
| '--ransac_times', type=int, default=None, help='repeat ransac multiple times for more robust evaluation') | |
| parser.add_argument( | |
| '--rmbd', type=int, default=None, help='remove border matches') | |
| parser.add_argument( | |
| '--deter', action='store_true', default=False, help='use deterministic mode for testing') | |
| parser.add_argument( | |
| '--half', action='store_true', default=False, help='pure16') | |
| parser.add_argument( | |
| '--flash', action='store_true', default=False, help='flash') | |
| parser = pl.Trainer.add_argparse_args(parser) | |
| return parser.parse_args() | |
| def inplace_relu(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('ReLU') != -1: | |
| m.inplace=True | |
| if __name__ == '__main__': | |
| # parse arguments | |
| args = parse_args() | |
| 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) | |
| if args.deter: | |
| torch.backends.cudnn.deterministic = True | |
| pl.seed_everything(config.TRAINER.SEED) # reproducibility | |
| # tune when testing | |
| if args.thr is not None: | |
| config.LOFTR.MATCH_COARSE.THR = args.thr | |
| if args.scannetX is not None and args.scannetY is not None: | |
| config.DATASET.SCAN_IMG_RESIZEX = args.scannetX | |
| config.DATASET.SCAN_IMG_RESIZEY = args.scannetY | |
| if args.megasize is not None: | |
| config.DATASET.MGDPT_IMG_RESIZE = args.megasize | |
| if args.npe: | |
| if config.LOFTR.COARSE.ROPE: | |
| assert config.DATASET.NPE_NAME is not None | |
| if config.DATASET.NPE_NAME is not None: | |
| if config.DATASET.NPE_NAME == 'megadepth': | |
| config.LOFTR.COARSE.NPE = [832, 832, config.DATASET.MGDPT_IMG_RESIZE, config.DATASET.MGDPT_IMG_RESIZE] # [832, 832, 1152, 1152] | |
| elif config.DATASET.NPE_NAME == 'scannet': | |
| config.LOFTR.COARSE.NPE = [832, 832, config.DATASET.SCAN_IMG_RESIZEX, config.DATASET.SCAN_IMG_RESIZEX] # [832, 832, 640, 640] | |
| else: | |
| config.LOFTR.COARSE.NPE = [832, 832, 832, 832] | |
| if args.ransac_times is not None: | |
| config.LOFTR.EVAL_TIMES = args.ransac_times | |
| if args.rmbd is not None: | |
| config.LOFTR.MATCH_COARSE.BORDER_RM = args.rmbd | |
| if args.pixel_thr is not None: | |
| config.TRAINER.RANSAC_PIXEL_THR = args.pixel_thr | |
| if args.ransac is not None: | |
| config.TRAINER.POSE_ESTIMATION_METHOD = args.ransac | |
| if args.ransac == 'LO-RANSAC' and config.TRAINER.RANSAC_PIXEL_THR == 0.5: | |
| config.TRAINER.RANSAC_PIXEL_THR = 2.0 | |
| if args.fp32: | |
| config.LOFTR.MP = False | |
| if args.half: | |
| config.LOFTR.HALF = True | |
| config.DATASET.FP16 = True | |
| else: | |
| config.LOFTR.HALF = False | |
| config.DATASET.FP16 = False | |
| if args.flash: | |
| config.LOFTR.COARSE.NO_FLASH = False | |
| loguru_logger.info(f"Args and config initialized!") | |
| # lightning module | |
| profiler = build_profiler(args.profiler_name) | |
| model = PL_LoFTR(config, pretrained_ckpt=args.ckpt_path, profiler=profiler, dump_dir=args.dump_dir) | |
| loguru_logger.info(f"LoFTR-lightning initialized!") | |
| # lightning data | |
| data_module = MultiSceneDataModule(args, config) | |
| loguru_logger.info(f"DataModule initialized!") | |
| # lightning trainer | |
| trainer = pl.Trainer.from_argparse_args(args, replace_sampler_ddp=False, logger=False) | |
| loguru_logger.info(f"Start testing!") | |
| trainer.test(model, datamodule=data_module, verbose=False) |