Enviroment setup:
conda env create -f environment.yaml
conda activate stfm
pip install torch einops yacs kornia
We provide training and test code for both indoor and outdoor datasets. The code scripts include training scripts for mono-modal baseline, multi-modal teacher model, and student-teacher learning model. The inference code is also provided. We also provide a simple demo code for performance visualization.
Both indoor and outdoor training code are released. To reproduce the result, the multi-modal teacher model need to be trained first and used for student-teacher learning.
conda activate stfm
bash ./scripts/train/indoor_ds_rgbd.sh ##teacher model training
## After teacher model training
bash ./scripts/train/indoor_ds_rgbd_t_s.sh ##student-teacher learning
conda activate stfm
bash ./scripts/test/indoor_ds.sh ##teacher model training
## if you want to see the multi-modal model's performance
bash ./scripts/test/indoor_ds_rgbd.sh
We also provide demo script for visualization. We provide the download link The link include both indoor and outdoor model weights. Please put the ckpt files in the folder weights.
conda activate stfm
# python script
python demo.py ##this sample code use indoor student model for correspondence prediction