This paper has been accepted by IEEE Transactions on Computational Imaging (TCI).
Project page | Paper | Pre-trained Models
Our synthetic data for the dynamic scene (videos) and static scene (image sequences) is available.
Our code was developed and tested on Ubuntu 20.04/CentOS 7 and Python 3.8.
First, clone our repo. Then,
cd code
pip install -r requirements. txt
First, train the tilt-removal model
For the dynamic scene modality, run the following:
python train_tilt_dynamic.py --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files}
Similarly, for the static scene modality, run the following:
python train_tilt_static.py --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files}
Next, train the deblurring model
For the dynamic scene modality, run the following:
python train_TMT_dynamic_2stage.py --path_tilt ${your_tilt_removal_model_path} --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files} --run_name ${your_exp_name}
Similarly, for the static scene modality, run the following:
python train_TMT_static_2stage.py --path_tilt ${your_tilt_removal_model_path} --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files} --run_name ${your_exp_name}
Alternatively, you can directly train a one-stage model
For the dynamic scene modality, run the following:
python train_TMT_dynamic.py --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files} --run_name ${your_exp_name}
Similarly, for the static scene modality, run the following:
python train_TMT_static.py --train_path ${your_training_data_path} --val_path ${your_validation_data_path} --log_path ${path_for_log_files} --run_name ${your_exp_name}
Test the two-stage model (tilt_removal + deblurring)
For the dynamic scene modality, run the following:
python test_TMT_dynamic_2stage.py --path_tilt ${your_tilt_removal_model_path} --model_path ${your_deblurring_model_path} --data_path ${your_validation_data_path} --result_path ${path_to_save_results}
Similarly, for the static scene modality, run the following:
python test_TMT_static_2stage.py --path_tilt ${your_tilt_removal_model_path} --model_path ${your_deblurring_model_path} --data_path ${your_validation_data_path} --result_path ${path_to_save_results}
Test the one-stage model
For the dynamic scene modality, run the following:
python test_TMT_dynamic.py --model_path ${your_model_path} --data_path ${your_validation_data_path} --result_path ${path_to_save_results}
Similarly, for the static scene modality, run the following:
python test_TMT_static.py --model_path ${your_model_path} --data_path ${your_validation_data_path} --result_path ${path_to_save_results}
All checkpoints saved during training are stored at ${log_path}/{run_name}/checkpoints. To customize your training pipeline, you can read and use the arguments in the Python scripts. The testing may take 10-20 hours.
If you find our work helps, please consider citing our work:
@misc{zhang2022TMT,
doi = {10.48550/ARXIV.2207.06465},
url = {https://arxiv.org/abs/2207.06465},
author = {Zhang, Xingguang and Mao, Zhiyuan and Chimitt, Nicholas and Chan, Stanley H.},
title = {Imaging through the Atmosphere using Turbulence Mitigation Transformer},
publisher = {arXiv},
year = {2022}
}