Home

Awesome

OverlapMamba: Novel Shift State Space Model for LiDAR-based Place Recognition

<p align="center"> <img src="./fig.png" width="50%"> </p> <b><p align="center" style="margin-top: -20px;"> Fig. 1 Core idea of the proposed OverlapMambaΒ </b></p>

πŸ’ΎEnvironment

We use pytorch-gpu for neural networks.

To use a GPU, first you need to install the nvidia driver and CUDA.

πŸ“–How to use

We provide a training and test tutorial for KITTI sequences in this repository. Before any operation, please modify the config file according to your setups.

πŸ“šDataset

Download KITTI dataset from KITTI.

We recommend you follow our code and data structures as follows.

Code Structure

β”œβ”€β”€ config
β”‚Β Β  β”œβ”€β”€ config_nclt.yml
β”‚Β Β  └── config.yml
β”œβ”€β”€ modules
β”‚Β Β  β”œβ”€β”€ loss.py
β”‚Β Β  β”œβ”€β”€ netvlad.py
β”‚Β Β  β”œβ”€β”€ overlap_mamba_nclt.py
β”‚Β Β  └── overlap_mamba.py
β”œβ”€β”€ test
β”‚Β Β  β”œβ”€β”€ test_mamba_topn_prepare.py
β”‚Β Β  β”œβ”€β”€ test_mamba_topn.py
β”‚Β Β  β”œβ”€β”€ test_kitti00_prepare.py
β”‚Β Β  β”œβ”€β”€ test_kitti00_PR.py
β”‚Β Β  β”œβ”€β”€ test_kitti00_topN.py
β”‚Β Β  β”œβ”€β”€ test_results_nclt
β”‚Β Β  β”‚Β Β  └── predicted_des_L2_dis_bet_traj_forward.npz (to be generated)
β”‚Β Β  └── test_results_kitti
β”‚Β Β      └── predicted_des_L2_dis.npz (to be generated)
β”œβ”€β”€ tools
β”‚Β Β  β”œβ”€β”€ read_all_sets.py
β”‚Β Β  β”œβ”€β”€ read_samples_haomo.py
β”‚Β Β  β”œβ”€β”€ read_samples.py
β”‚Β Β  └── utils
β”‚Β Β      β”œβ”€β”€ gen_depth_data.py
β”‚Β Β      β”œβ”€β”€ split_train_val.py
β”‚Β Β      └── utils.py
β”œβ”€β”€ train
β”‚Β Β  β”œβ”€β”€ training_overlap_mamba_nclt.py
β”‚Β Β  └── training_overlap_mamba_kitti.py
β”œβ”€β”€ valid
β”‚Β Β  └── valid_seq.py
β”œβ”€β”€ visualize
β”‚Β Β  β”œβ”€β”€ des_list.npy
β”‚Β Β  └── viz_nclt.py
└── weights
    β”œβ”€β”€ pretrained_overlap_mamba_nclt.pth.tar
    └── pretrained_overlap_mamba.pth.tar

Dataset Structure

  data_root_folder (KITTI sequences root) follows:
  β”œβ”€β”€ 00
  β”‚Β Β  β”œβ”€β”€ depth_map
  β”‚Β Β    β”œβ”€β”€ 000000.png
  β”‚Β Β    β”œβ”€β”€ 000001.png
  β”‚Β Β    β”œβ”€β”€ 000002.png
  β”‚Β Β    β”œβ”€β”€ ...
  β”‚Β Β  └── overlaps
  β”‚Β Β    β”œβ”€β”€ train_set.npz
  β”œβ”€β”€ 01
  β”œβ”€β”€ 02
  β”œβ”€β”€ ...
  β”œβ”€β”€ 10
  └── loop_gt_seq00_0.3overlap_inactive.npz
  
  valid_scan_folder (KITTI sequence 02 velodyne) contains:
  β”œβ”€β”€ 000000.bin
  β”œβ”€β”€ 000001.bin
  ...

  gt_valid_folder (KITTI sequence 02 computed overlaps) contains:
  β”œβ”€β”€ 02
  β”‚Β Β  β”œβ”€β”€ overlap_0.npy
  β”‚Β Β  β”œβ”€β”€ overlap_10.npy
  ...

You need to download or generate the following files and put them in the right positions of the structure above:

Quick Use

For a quick use, you could download our model pretrained on KITTI, and the following two files also should be downloaded :

Then you should modify demo1_config in the file config.yaml.

Run the demo by:

cd demo
python ./demo_compute_overlap_sim.py

You can see a query scan (000000.bin of KITTI 00) with a reprojected positive sample (000005.bin of KITTI 00) and a reprojected negative sample (000015.bin of KITTI 00), and the corresponding similarity.

Training

In the file config.yaml, training_seqs are set for the KITTI sequences used for training.

You can start the training with

cd train
python ./training_overlap_mamba_kitti.py

You can resume from our pretrained model here for training.

Testing

Once a model has been trained , the performance of the network can be evaluated. Before testing, the parameters shoud be set in config.yaml

Therefore you can start the testing scripts as follows:

cd test
mkdir test_results_kitti
python test_kitti00_prepare.py
python test_kitti00_PR.py
python test_kitti00_topN.py

After you run test_kitti00_prepare.py, a file named predicted_des_L2_dis.npz is generated in test_results_kitti, which is used by python test_kitti00_PR.py to calculate PR curve and F1max, and used by python test_kitti00_topN.py to calculate topN recall.

Performance

The code of this repo uses a minimal PyTorch implementation of Mamba, which implements the scan operation as a sequential loop. It closely follows the file of the official Mamba implementation, but its performance are a bit worse than the official Mamba implementation.

After our paper is accepted, we will provide a C++ implementation of OverlapMamba with libtorch for faster retrival.

πŸ‘Acknowledgment

This repo is based on OverlapTransformer and mamba.py, we are very grateful for their excellent work and appreciate their contributions to LiDAR-based place recognition(LPR) and highly recommend people to use their excellent public available code.