pp00704831 / BANet-TIP-2022

37 stars 2 forks source link

BANet (TIP 2022)

Pytorch Implementation of "BANet: Blur-aware Attention Networks for Dynamic Scene Deblurring"

Installation

The implementation of our BANet is modified from "DeblurGANv2"

git clone https://github.com/pp00704831/BANet.git
cd BANet
conda create -n banet python=3.6
source activate banet
conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch
pip install opencv-python tqdm pyyaml joblib glog scikit-image tensorboardX albumentations
pip install -U albumentations[imgaug]
pip install albumentations==1.1.0

Training

Download "GoPro" dataset into './datasets'
For example: './datasets/GoPro/train/blur/**/*.png'

We train our BANner in two stages:
1) We pre-train for 3000 epochs on patch size 256x256. Please run the following commands.

python pretrained.py

2) We fine-tune for 1000 epochs on patch size 512x512. Please run the following commands.

python train.py

Testing

Download "GoPro" dataset into './datasets'
For example: './datasets/GoPro/test/blur/**/*.png'
For reproducing our results, download the GoPro trained model "BANet_GoPro.pth

For testing on GoPro dataset

python predict_BANet_GoPro_test_results.py --weights_path ./BANet_GoPro.pth 

For testing on HIDE dataset
Download "HIDE" dataset into './datasets'

python predict_BANet_HIDE_test_results.py --weights_path ./BANet_GoPro.pth 

For testing your own training weight on GoPro or HIDE

Take GoPro for example

Evaluation

For evaluation on GoPro results in MATLAB Download "BANet_GoPro_result" into './out'

evaluation_GoPro.m

For evaluation on HIDE results in MATLAB Download "BANet_HIDE_result" into './out'

evaluation_HIDE.m

BANet+

The training strategy is the same as BANet above.

Testing

Download the GoPro trained model "BANet_Plus_GoPro.pth
For testing on GoPro dataset

python predict_BANet_Plus_GoPro_test_results.py --weights_path ./BANet_Plus_GoPro.pth 

For testing on HIDE dataset

python predict_BANet_Plus_HIDE_test_results.py --weights_path ./BANet_Plus_GoPro.pth 

Citation

@ARTICLE{Tsai2022BANet,
  author={Tsai, Fu-Jen and Peng, Yan-Tsung and Tsai, Chung-Chi and Lin, Yen-Yu and Lin, Chia-Wen},
  journal={IEEE Transactions on Image Processing}, 
  title={BANet: A Blur-Aware Attention Network for Dynamic Scene Deblurring}, 
  year={2022},
  volume={31},
  pages={6789-6799},
}