ifitsmanu / TMT

Project page of the Turbulence Mitigation Transformer

Turbulence Mitigation Transformer

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.

Quick Start

Installation

First, clone our repo. Then,

cd code
pip install -r requirements. txt

Training

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}

Testing

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}
}