🔥 Project Page | Paper

🧩 Dataset and Pre-train Models

🔑 Setup and Prepare LMDB files

cd code
pip install -r requirements.txt

LMDB format is encouraged for our dynamic scene data. Before preparing your LMDB file from the downloaded ATSyn-dynamic dataset, please ensure you have 2TB of free space on your server. You will need to modify the path variables in the make_lmdb.py file. After that, just run

cd code
python make_lmdb.py

Alternatively, you can also use the .mp4 files for training, please refer to the code in data/dataset_video_train.py and training script in TMT for more information.

🛠️ Training

For the training on dynamic scene data, run the following:

python train_DATUM_dynamic.py --train_path ${your_training_data_path} --train_info ${the associated train_info.json} --val_path ${your_testing_data_path} --val_info ${the associated test_info.json} -f ${loaded_model_path (if you want to resume a pre-trained checkpoint)} 

Other arguments for training are described in the train_DATUM_dynamic.py file, please refer to them for more flexible training. Use smaller patch_size and num_frames in the beginning phase of training can accelerate the entire process.

Later, you can start finetuning on the static scene images for the static scene model by running the following:

python train_DATUM_static.py --train_path ${your_training_data_path} --val_path ${your_testing_data_path} -f ${pretrained_dynamic_scene_model_path} 

We injected a certain level of Gaussian noise during training in both modalities for better generalization on real-world data.

🚀 Performance Evaluation

Dynamic scene model on ATSyn_dynamic dataset:

python test_DATUM_dynamic.py --data_path ${your_testing_data_path} --info_path ${the associated test_info.json} -result ${path_for_stored_output} -mp ${testing_model_path} 

Static scene model on ATSyn_static dataset:

python test_DATUM_static.py --val_path ${your_testing_data_path} -result ${path_for_stored_output} -f ${testing_model_path} 

Inference on Turbulence Text dataset, we generate the central 4 frames for the text recognition:

python inference_DATUM_text.py -f ${testing_static_scene_model_path} --n_frames 60 --resize 360

Please modify the path of the input and output images in the inference_DATUM_text.py

Please refer to CWSSIM_static.py and CWSSIM_dynamic.py to evaluate the CW-SSIM score.

👍 Useful Links

Turbulence @ Purdue i2Lab | UG2 workshop @ CVPR 2024

Restoration:

Link Zhang, Xingguang, Zhiyuan Mao, Nicholas Chimitt, and Stanley H. Chan. "Imaging through the atmosphere using turbulence mitigation transformer." IEEE Transactions on Computational Imaging (2024).

Link Ajay Jaiswal, Xingguang Zhang, Stanley H. Chan, Zhangyang Wang. "Physics-Driven Turbulence Image Restoration with Stochastic Refinement." Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2023, pp. 12170-12181

Link Liang, Jingyun, Yuchen Fan, Xiaoyu Xiang, Rakesh Ranjan, Eddy Ilg, Simon Green, Jiezhang Cao, Kai Zhang, Radu Timofte, and Luc V. Gool. "Recurrent video restoration transformer with guided deformable attention." Advances in Neural Information Processing Systems 35 (2022): 378-393.

Link Jin, Darui, Ying Chen, Yi Lu, Junzhang Chen, Peng Wang, Zichao Liu, Sheng Guo, and Xiangzhi Bai. "Neutralizing the impact of atmospheric turbulence on complex scene imaging via deep learning." Nature Machine Intelligence 3, no. 10 (2021): 876-884.

Datasets:

OTIS dataset | TSRWGAN data | Turbulence Text | Heat Chamber | TMT dataset | BRIAR (Not public yet)

📘 Citation

Please consider citing our work as follows if it is helpful.

@InProceedings{zhang2024spatio,
    author={Zhang, Xingguang and Chimitt, Nicholas and Chi, Yiheng and Mao, Zhiyuan and Chan, Stanley H}, 
    title={Spatio-Temporal Turbulence Mitigation: A Translational Perspective},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month={June},
    year={2024}
}