ADD LoFTR and some examples

.gitignore

@@ -0,0 +1,12 @@
+.vscode/
+*.pyc
+*.DS_Store
+*.swp
+*.pth
+tmp.*
+*/.ipynb_checkpoints/*
+
+logs/
+weights/
+dump/
+src/loftr/utils/superglue.py

README.md

@@ -9,10 +9,82 @@
 ![demo_vid](assets/loftr-github-demo.gif)
 
 
+## Installation
+```shell
+# For full pytorch-lightning trainer features
+conda env create -f environment.yaml
+conda activate loftr
+
+# For the LoFTR matcher only
+pip install torch einops yacs kornia
+```
+
+We provide the [download link](https://drive.google.com/drive/folders/1DOcOPZb3-5cWxLqn256AhwUVjBPifhuf?usp=sharing) to 
+  - the scannet-1500-testset (~1GB).
+  - the megadepth-1500-testset (~600MB).
+  - 4 pretrained models of indoor-ds, indoor-ot, outdoor-ds and outdoor-ot (each ~45MB).
+
+By now, the LoFTR-DS model is ready to go!
+
+<details>
+  <summary>[Requirements for LoFTR-OT]</summary>
+
+  We use the code from [SuperGluePretrainedNetwork](https://github.com/magicleap/SuperGluePretrainedNetwork) for optimal transport. However, we can't provide the code directly due to its LICENSE. We recommend downloading it instead. 
+
+  ```shell
+  cd src/loftr/utils  
+  wget https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/master/models/superglue.py 
+  ```
+</details>
+
+
+## Run the code
+
+### Match image pairs with LoFTR
+
+<details>
+  <summary>[code snippets]</summary>
+
+  ```python
+  from src.loftr import LoFTR, default_cfg
+
+  # Initialize LoFTR
+  matcher = LoFTR(config=default_cfg)
+  matcher.load_state_dict(torch.load("weights/indoor_ds.ckpt")['state_dict'])
+  matcher = matcher.eval().cuda()
+
+  # Inference
+  with torch.no_grad():
+      matcher(batch)    # batch = {'image0': img0, 'image1': img1}
+      mkpts0 = batch['mkpts0_f'].cpu().numpy()
+      mkpts1 = batch['mkpts1_f'].cpu().numpy()
+  ```
+
+</details>
+
+An example is in the `notebooks/demo_single_pair.ipynb`.
+
+### Reproduce the testing results with pytorch-lightning
+
+```shell
+conda activate loftr
+# with shell script
+bash ./scripts/reproduce_test/indoor_ds.sh
+
+# or
+python test.py configs/data/scannet_test_1500.py configs/loftr/loftr_ds.py --ckpt_path weights/indoor_ds.ckpt --profiler_name inference --gpus=1 --accelerator="ddp"
+```
+
+For visualizing the dump results, please refer to `notebooks/visualize_dump_results.ipynb`.
+
+### Reproduce the training phase with pytorch-lightning
+
+The code is coming soon, stay tuned!
+
 <br/>
 
+
 ## Code release ETA
-We plan to release the inference-only code and pretrained model within the upcoming week, stay tuned. 
 The entire codebase for data pre-processing, training and validation is under major refactoring and will be released around June.
 Please subscribe to [this discussion thread](https://github.com/zju3dv/LoFTR/discussions/2) if you wish to be notified of the code release.
 In the meanwhile, discussions about the paper are welcomed in the [discussion panel](https://github.com/zju3dv/LoFTR/discussions).

assets/megadepth_test_1500_scene_info/megadepth_test_1500.txt

@@ -0,0 +1,5 @@
+0022_0.1_0.3.npz
+0015_0.1_0.3.npz
+0015_0.3_0.5.npz
+0022_0.3_0.5.npz
+0022_0.5_0.7.npz

assets/scannet_test_1500/scannet_test.txt

@@ -0,0 +1 @@
+test.npz

assets/scannet_test_1500/statistics.json

@@ -0,0 +1,102 @@
+{
+    "scene0707_00": 15,
+    "scene0708_00": 15,
+    "scene0709_00": 15,
+    "scene0710_00": 15,
+    "scene0711_00": 15,
+    "scene0712_00": 15,
+    "scene0713_00": 15,
+    "scene0714_00": 15,
+    "scene0715_00": 15,
+    "scene0716_00": 15,
+    "scene0717_00": 15,
+    "scene0718_00": 15,
+    "scene0719_00": 15,
+    "scene0720_00": 15,
+    "scene0721_00": 15,
+    "scene0722_00": 15,
+    "scene0723_00": 15,
+    "scene0724_00": 15,
+    "scene0725_00": 15,
+    "scene0726_00": 15,
+    "scene0727_00": 15,
+    "scene0728_00": 15,
+    "scene0729_00": 15,
+    "scene0730_00": 15,
+    "scene0731_00": 15,
+    "scene0732_00": 15,
+    "scene0733_00": 15,
+    "scene0734_00": 15,
+    "scene0735_00": 15,
+    "scene0736_00": 15,
+    "scene0737_00": 15,
+    "scene0738_00": 15,
+    "scene0739_00": 15,
+    "scene0740_00": 15,
+    "scene0741_00": 15,
+    "scene0742_00": 15,
+    "scene0743_00": 15,
+    "scene0744_00": 15,
+    "scene0745_00": 15,
+    "scene0746_00": 15,
+    "scene0747_00": 15,
+    "scene0748_00": 15,
+    "scene0749_00": 15,
+    "scene0750_00": 15,
+    "scene0751_00": 15,
+    "scene0752_00": 15,
+    "scene0753_00": 15,
+    "scene0754_00": 15,
+    "scene0755_00": 15,
+    "scene0756_00": 15,
+    "scene0757_00": 15,
+    "scene0758_00": 15,
+    "scene0759_00": 15,
+    "scene0760_00": 15,
+    "scene0761_00": 15,
+    "scene0762_00": 15,
+    "scene0763_00": 15,
+    "scene0764_00": 15,
+    "scene0765_00": 15,
+    "scene0766_00": 15,
+    "scene0767_00": 15,
+    "scene0768_00": 15,
+    "scene0769_00": 15,
+    "scene0770_00": 15,
+    "scene0771_00": 15,
+    "scene0772_00": 15,
+    "scene0773_00": 15,
+    "scene0774_00": 15,
+    "scene0775_00": 15,
+    "scene0776_00": 15,
+    "scene0777_00": 15,
+    "scene0778_00": 15,
+    "scene0779_00": 15,
+    "scene0780_00": 15,
+    "scene0781_00": 15,
+    "scene0782_00": 15,
+    "scene0783_00": 15,
+    "scene0784_00": 15,
+    "scene0785_00": 15,
+    "scene0786_00": 15,
+    "scene0787_00": 15,
+    "scene0788_00": 15,
+    "scene0789_00": 15,
+    "scene0790_00": 15,
+    "scene0791_00": 15,
+    "scene0792_00": 15,
+    "scene0793_00": 15,
+    "scene0794_00": 15,
+    "scene0795_00": 15,
+    "scene0796_00": 15,
+    "scene0797_00": 15,
+    "scene0798_00": 15,
+    "scene0799_00": 15,
+    "scene0800_00": 15,
+    "scene0801_00": 15,
+    "scene0802_00": 15,
+    "scene0803_00": 15,
+    "scene0804_00": 15,
+    "scene0805_00": 15,
+    "scene0806_00": 15
+}

configs/data/base.py

@@ -0,0 +1,31 @@
+"""
+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_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_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_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = None
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# dataset config
+_CN.DATASET.MIN_OVERLAP_SCORE = 0.4
+
+cfg = _CN

configs/data/megadepth_test_1500.py

@@ -0,0 +1,11 @@
+from configs.data.base import cfg
+
+TEST_BASE_PATH = "assets/megadepth_test_1500_scene_info"
+
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TEST_DATA_ROOT = "/data/MegaDepth/megadepth_test_1500"
+cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}"
+cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/megadepth_test_1500.txt"
+
+cfg.DATASET.MGDPT_IMG_RESIZE = 840
+cfg.DATASET.MIN_OVERLAP_SCORE = 0.0

configs/data/scannet_test_1500.py

@@ -0,0 +1,11 @@
+from configs.data.base import cfg
+
+TEST_BASE_PATH = "assets/scannet_test_1500"
+
+cfg.DATASET.TEST_DATA_SOURCE = "ScanNet"
+cfg.DATASET.TEST_DATA_ROOT = "/data/scannet/scannet_test_1500"
+cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}"
+cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/scannet_test.txt"
+cfg.DATASET.TEST_INTRINSIC_PATH = f"{TEST_BASE_PATH}/intrinsics.npz"
+
+cfg.DATASET.MIN_OVERLAP_SCORE = 0.0

configs/loftr/loftr_ds.py

@@ -0,0 +1,3 @@
+from src.config.default import _CN as cfg
+
+cfg.LOFTR.MATCH_COARSE.MATCH_TYPE = 'dual_softmax'

configs/loftr/loftr_ot.py

@@ -0,0 +1,3 @@
+from src.config.default import _CN as cfg
+
+cfg.LOFTR.MATCH_COARSE.MATCH_TYPE = 'sinkhorn'

environment.yaml

@@ -0,0 +1,14 @@
+name: loftr
+channels:
+  # - https://dx-mirrors.sensetime.com/anaconda/cloud/pytorch
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.8
+  - cudatoolkit=10.2
+  - pytorch=1.8.0
+  - pytorch-lightning<=1.1.8  # https://github.com/PyTorchLightning/pytorch-lightning/issues/6318
+  - pip
+  - pip:
+      - -r file:requirements.txt

requirements.txt

@@ -0,0 +1,14 @@
+opencv_python==4.4.0.46
+albumentations==0.5.1 --no-binary=imgaug,albumentations
+ray>=1.0.1
+einops==0.3.0
+kornia==0.4.1
+loguru==0.5.3
+yacs>=0.1.8
+tqdm
+autopep8
+pylint
+ipython
+jupyterlab
+matplotlib
+h5py==3.1.0

scripts/reproduce_test/indoor_ds.sh

@@ -0,0 +1,29 @@
+#!/bin/bash -l
+
+SCRIPTPATH=$(dirname $(readlink -f "$0"))
+PROJECT_DIR="${SCRIPTPATH}/../../"
+
+# conda activate loftr
+export PYTHONPATH=$PROJECT_DIR:$PYTHONPATH
+cd $PROJECT_DIR
+
+data_cfg_path="configs/data/scannet_test_1500.py"
+main_cfg_path="configs/loftr/loftr_ds.py"
+ckpt_path="weights/indoor_ds.ckpt"
+dump_dir="dump/loftr_ds_indoor"
+profiler_name="inference"
+n_nodes=1  # mannually keep this the same with --nodes
+n_gpus_per_node=-1
+torch_num_workers=4
+batch_size=1  # per gpu
+
+python -u ./test.py \
+    ${data_cfg_path} \
+    ${main_cfg_path} \
+    --ckpt_path=${ckpt_path} \
+    --dump_dir=${dump_dir} \
+    --gpus=${n_gpus_per_node} --num_nodes=${n_nodes} --accelerator="ddp" \
+    --batch_size=${batch_size} --num_workers=${torch_num_workers}\
+    --profiler_name=${profiler_name} \
+    --benchmark 
+

scripts/reproduce_test/indoor_ot.sh

@@ -0,0 +1,29 @@
+#!/bin/bash -l
+
+SCRIPTPATH=$(dirname $(readlink -f "$0"))
+PROJECT_DIR="${SCRIPTPATH}/../../"
+
+# conda activate loftr
+export PYTHONPATH=$PROJECT_DIR:$PYTHONPATH
+cd $PROJECT_DIR
+
+data_cfg_path="configs/data/scannet_test_1500.py"
+main_cfg_path="configs/loftr/loftr_ot.py"
+ckpt_path="weights/indoor_ot.ckpt"
+dump_dir="dump/loftr_ot_indoor"
+profiler_name="inference"
+n_nodes=1  # mannually keep this the same with --nodes
+n_gpus_per_node=-1
+torch_num_workers=4
+batch_size=1  # per gpu
+
+python -u ./test.py \
+    ${data_cfg_path} \
+    ${main_cfg_path} \
+    --ckpt_path=${ckpt_path} \
+    --dump_dir=${dump_dir} \
+    --gpus=${n_gpus_per_node} --num_nodes=${n_nodes} --accelerator="ddp" \
+    --batch_size=${batch_size} --num_workers=${torch_num_workers}\
+    --profiler_name=${profiler_name} \
+    --benchmark 
+

scripts/reproduce_test/outdoor_ds.sh

@@ -0,0 +1,29 @@
+#!/bin/bash -l
+
+SCRIPTPATH=$(dirname $(readlink -f "$0"))
+PROJECT_DIR="${SCRIPTPATH}/../../"
+
+# conda activate loftr
+export PYTHONPATH=$PROJECT_DIR:$PYTHONPATH
+cd $PROJECT_DIR
+
+data_cfg_path="configs/data/megadepth_test_1500.py"
+main_cfg_path="configs/loftr/loftr_ds.py"
+ckpt_path="weights/outdoor_ds.ckpt"
+dump_dir="dump/loftr_ds_outdoor"
+profiler_name="inference"
+n_nodes=1  # mannually keep this the same with --nodes
+n_gpus_per_node=-1
+torch_num_workers=4
+batch_size=1  # per gpu
+
+python -u ./test.py \
+    ${data_cfg_path} \
+    ${main_cfg_path} \
+    --ckpt_path=${ckpt_path} \
+    --dump_dir=${dump_dir} \
+    --gpus=${n_gpus_per_node} --num_nodes=${n_nodes} --accelerator="ddp" \
+    --batch_size=${batch_size} --num_workers=${torch_num_workers}\
+    --profiler_name=${profiler_name} \
+    --benchmark 
+