huicongzhang / bsstnet

Implementation of "Blur-aware Spatio-temporal Sparse Transformer for Video Deblurring". (Zhang et al., CVPR 2024)

Blur-aware Spatio-temporal Sparse Transformer for Video Deblurring

Huicong Zhang, Haozhe Xie, Hongxun Yao

Harbin Institute of Technology, S-Lab, Nanyang Technological University

Update

• [2024/07/04] The training and testing code are released.
• [2024/02/29] The repo is created.

Datasets

We use the GoPro and DVD datasets in our experiments, which are available below:

• GoPro
• DVD

You could download the zip file and then extract it to the datasets folder.

Pretrained Models

You could download the pretrained model from here and put the weights in model_zoos.

Prerequisites

Clone the Code Repository

git clone https://github.com/huicongzhang/BSSTNet.git

Install Denpendencies

conda create -n BSSTNet python=3.8
conda activate BSSTNet
pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install mmcv-full==1.7.1 -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9/index.html
pip install -r requirements.txt
BASICSR_EXT=True python setup.py develop

Test

To train BSSTNet, you can simply use the following commands:

GoPro dataset

scripts/dist_test.sh 2 options/test/BSST/gopro_BSST.yml

DVD dataset

scripts/dist_test.sh 2 options/test/BSST/dvd_BSST.yml

Train

To train BSSTNet, you can simply use the following commands:

GoPro dataset

scripts/dist_train.sh 2 options/test/BSST/gopro_BSST.yml

DVD dataset

scripts/dist_train.sh 2 options/test/BSST/dvd_BSST.yml

Cite this work

@inproceedings{zhang2024bsstnet,
  title     = {Blur-aware Spatio-temporal Sparse Transformer for Video Deblurring},
  author    = {Zhang, Huicong and 
               Xie, Haozhe and 
               Yao, Hongxun},
  booktitle = {CVPR},
  year      = {2024}
}

License

This project is open sourced under MIT license.

Acknowledgement

This project is based on BasicSR, ProPainter and Shift-Net.

bsstnet's Issues

Questions about downloading your paper

Hello,

Congratulations on your CVPR paper! Do you plan to release it on ArXiv, or could I possibly download it recently? Thank you very much.

get_bi_flows error

Hi, could you test the get_bi_flows snippet in Video_BSST_model.py.

I tested it with following:

import torch
lqs = torch.rand(1, 10, 3, 180, 320).cuda()
flows_forward, flows_backward = get_bi_flows(lqs)

But it turns error in Line 66 that
Exception has occurred: RuntimeError grid_sampler(): expected grid and input to have same batch size, but got input with sizes [2160, 1, 12, 20] and grid with sizes [1980, 9, 9, 2]

Thanks!

About optical flow model

Hi, thanks for you excellent work! In your paper, when reporting runtime or parameters, did you include the computation of these parameters about RAFT? I think the computational cost of RAFT should be included.