lu-feng / dhe-vpr

Official repository for the AAAI 2024 paper "Deep Homography Estimation for Visual Place Recognition".

DHE-VPR

This is the official repository for the AAAI 2024 paper "Deep Homography Estimation for Visual Place Recognition". [AAAI proceedings] [arXiv] The arXiv version is more complete.

Our another two-stage VPR work SelaVPR achieved SOTA performance on several datasets. The code has been released HERE.

Getting Started

This repo follows the Visual Geo-localization Benchmark. You can refer to it (VPR-datasets-downloader) to prepare datasets and train the CCT-14 backbone (i.e. feature extractor).

The dataset should be organized in a directory tree as such:

├── datasets_vg
    └── datasets
        └── pitts30k
            └── images
                ├── train
                │   ├── database
                │   └── queries
                ├── val
                │   ├── database
                │   └── queries
                └── test
                    ├── database
                    └── queries

You can directly download the trained CCT-14 backbone:

trained on MSLS: CCT14_msls

trained on Pitts30k: CCT14_pitts30k

Train (initialize) the DHE network

After getting the CCT14 backbone trained on MSLS (CCT14_msls.pth), you can train (i.e. initialize) the DHE network on MSLS:

python train_dhe.py --resume_fe=/path/to/your/CCT14_msls.pth --datasets_folder=/path/to/your/datasets_vg/datasets --dataset_name=msls

You can directly download the initialized DHE network HERE.

Finetune

To jointly finetune the backbone and the DHE network on the MSLS dataset, please run:

python3 finetune.py --datasets_folder=/path/to/your/datasets_vg/datasets --dataset_name=msls --epochs_num=2 --resume_fe=/path/to/your/CCT14_msls.pth --resume_hr=/path/to/your/initializedDHE.torch --queries_per_epoch=10000

Finetune on the Pitts30k dataset, please run:

python3 finetune.py --datasets_folder=/path/to/your/datasets_vg/datasets --dataset_name=pitts30k --epochs_num=40 --resume_fe=/path/to/your/CCT14_pitts30k.pth --resume_hr=/path/to/your/initializedDHE.torch

You can directly download the finetuned CCT14 backbone and DHE network:

MSLS: finetunedCCT14 | finetunedDHE

Pitts30k: finetunedCCT14 | finetunedDHE

Test

To evaluate the finetuned complete DHE-VPR model on MSLS (or Pitts30k), run:

python eval.py  --resume_fe=/path/to/your/finetunedCCT14_msls.torch --resume_hr=/path/to/your/finetunedDHE_msls.torch --datasets_folder=/path/to/your/datasets_vg/datasets --dataset_name=msls

Acknowledgements

Parts of this repo are inspired by the following repositories:

Visual Geo-localization Benchmark

GeoWarp

TransVPR

Citation

If you find this repo useful for your research, please consider citing the paper

@inproceedings{dhevpr,
  title={Deep Homography Estimation for Visual Place Recognition},
  author={Lu, Feng and Dong, Shuting and Zhang, Lijun and Liu, Bingxi and Lan, Xiangyuan and Jiang, Dongmei and Yuan, Chun},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  year={2024},
  volume={38}, 
  number={9}, 
  pages={10341-10349}
}

dhe-vpr's Issues

CUDA error: out of memory

Hello,
I'm currently in the process of attempting to replicate the remarkable work you've shared. However, I've encountered a hurdle along the way. Following the training pipeline outlined in the repository and adhering to the provided requirements, I attempted to train the DHE-VPR using 4090 with 24GB graphics memory, same size with the 3090 . Unfortunately, after training for 1 epoch, the program crashed and reported the following issue:

Traceback (most recent call last):
File "/home/xx/Documents/codespace/DHE-VPR/train_dhe.py", line 151, in
REIloss = homography_project.reprojection_error_ofinliers(model, queries_fw, positives_fw, weights=random_weights)
File "/home/xx/Documents/codespace/DHE-VPR/homography_project.py", line 102, in reprojection_error_ofinliers
reproject_error[i] = match_batch_tensor(query, pred, theta, trainflag=True, img_size=(384,384))
File "/home/xx/Documents/codespace/DHE-VPR/homography_project.py", line 31, in match_batch_tensor
max1 = torch.argmax(M, dim=1) #(N,l)
RuntimeError: CUDA error: out of memory

I believe this issue is related to memory allocation on the CUDA device. Your expertise and guidance in resolving this problem would be immensely valuable to me

Thank you for your time

question code release

Hello
Thanks for your interesting research!
Can you give me an rough dates the code will be released?