TAViT

This repository is the official implementation of "Task-Agnostic Vision Transformer for Distributed Learning of Image Processing" which is published in IEEE Transactions on Image Processing.

[TIP paper]

Requirements

To install requirements:

conda env create -f environment.yaml

Evaluation

📋 The inference in this code is performed by stacking all patches of an input image in batch dimension, while we report results by testing the patches with batch size 1.

To evaluate our model on deblocking task with quality 10, run:

python deblocking_inference.py --Q=10

To evaluate our model on denoising task with σ=30, run:

python denoising_inference.py --mean=0 --std=30

To evaluate our model on deraining task, run:

python deraining_inference.py 

To evaluate our model on deblurring task, run:

python deblurring_inference.py  

Pre-trained Models

You can download pre-trained model checkpoints (.pt) from the below link. https://drive.google.com/drive/folders/1j_EXwzkMhl2y0TCsOnwg68lbYL0Q07RN?usp=sharing.

Copy downloaded checkpoint file to proper directory under 'inference/checkpoints'.

Results (three cycles)

Deblocking

Quality	Metric	Input	DnCNN	AR-CNN	QCN	Base	Single	Multi
10	PSNR SSIM	25.67 0.719	26.70 0.755	26.42 0.777	27.66 0.811	27.67 0.785	27.59 0.785	27.69 0.786
50	PSNR SSIM	31.51 0.902	32.70 0.918	N/A N/A	33.00 0.934	33.01 0.923	32.93 0.924	33.20 0.924

Denoising

Metric	Input	CBM3D	FFDNet	IRCNN	DHDN	SADNet	Base	Single	Multi
PSNR SSIM	19.03 0.336	29.71 0.843	30.32 0.860	30.31 0.860	30.41 0.864	30.64 N/A	30.43 0.864	30.65 0.870	30.69 0.871

Deraining

Dataset	Metric	Input	DerainNet	SEMI	UMRL	PreNet	MSPFN	Base	Single	Multi
Rain100H	PSNR SSIM	12.13 0.349	14.92 0.592	16.56 0.486	26.01 0.832	26.77 0.858	28.66 0.860	28.88 0.863	28.95 0.864	29.35 0.875
Rain100L	PSNR SSIM	25.52 0.825	27.03 0.884	25.03 0.842	29.18 0.923	32.44 0.950	32.40 0.933	32.93 0.937	32.50 0.935	34.30 0.949

Deblurring

Metric	Input	DeblurGAN	Nah et al.	Zhang et al.	DeblurGANv2	Base	Single	Multi
PSNR SSIM	25.64 0.790	28.70 0.858	29.08 0.914	29.19 0.931	29.55 0.934	28.62 0.864	29.28 0.877	30.06 0.894

License

• We implemented out code based on Flower (https://github.com/adap/flower) which is under Apache License, Version 2.0.

• See LICENSE for more information.

Data source

Data for each task can be downloaded at following links:

• PASCAL VOC 2012 : http://host.robots.ox.ac.uk/pascal/VOC/

• BSDS500 : https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/resources.html

• Deraining :

  • Rain1400 : https://xueyangfu.github.io/projects/cvpr2017.html
  • Rain1800 : https://www.icst.pku.edu.cn/struct/Projects/joint\_rain\_removal.html
  • Rain800 : https://github.com/hezhangsprinter/ID-CGAN
  • Rain12 : https://yu-li.github.io/
  • Rain100L/Rain100H : https://www.icst.pku.edu.cn/struct/Projects/joint_rain_removal.html

• GoPro : https://seungjunnah.github.io/Datasets/gopro.html