MAXIM: Multi-Axis MLP for Image Processing (CVPR 2022 Oral)

This repo is the official implementation of [CVPR 2022 Oral] paper: "MAXIM: Multi-Axis MLP for Image Processing" by Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li

Google Research, University of Texas at Austin

Disclaimer: This is not an officially supported Google product.

---

Abstract: Recent progress on Transformers and multi-layer perceptron (MLP) models provide new network architectural designs for computer vision tasks. Although these models proved to be effective in many vision tasks such as image recognition, there remain challenges in adapting them for low-level vision. The inflexibility to support high-resolution images and limitations of local attention are perhaps the main bottlenecks. In this work, we present a multi-axis MLP based architecture called MAXIM, that can serve as an efficient and flexible general-purpose vision backbone for image processing tasks. MAXIM uses a UNet-shaped hierarchical structure and supports long-range interactions enabled by spatially-gated MLPs. Specifically, MAXIM contains two MLP-based building blocks: a multi-axis gated MLP that allows for efficient and scalable spatial mixing of local and global visual cues, and a cross-gating block, an alternative to cross-attention, which accounts for cross-feature conditioning. Both these modules are exclusively based on MLPs, but also benefit from being both global and `fully-convolutional', two properties that are desirable for image processing. Our extensive experimental results show that the proposed MAXIM model achieves state-of-the-art performance on more than ten benchmarks across a range of image processing tasks, including denoising, deblurring, deraining, dehazing, and enhancement while requiring fewer or comparable numbers of parameters and FLOPs than competitive models.

---

Architecture

Installation

Install dependencies:

pip install -r requirements.txt

Setup project:

pip install .

Results and Pre-trained models

We provide the pre-trained models and visual results. Please contact us if you have any questions or requests.

Task	Dataset	PSNR	SSIM	Model	#params	FLOPs	ckpt	outputs
Denoising	SIDD	39.96	0.960	MAXIM-3S	22.2M	339G	ckpt	images
Denoising	DND	39.84	0.954	MAXIM-3S	22.2M	339G	ckpt	images
Deblurring	GoPro	32.86	0.961	MAXIM-3S	22.2M	339G	ckpt	images
Deblurring	HIDE	32.83	0.956	MAXIM-3S	22.2M	339G	ckpt	images
Deblurring	REDS	28.93	0.865	MAXIM-3S	22.2M	339G	ckpt	images
Deblurring	RealBlur-R	39.45	0.962	MAXIM-3S	22.2M	339G	ckpt	images
Deblurring	RealBlur-J	32.84	0.935	MAXIM-3S	22.2M	339G	ckpt	images
Deraining	Rain13k	33.24	0.933	MAXIM-2S	14.1M	216G	ckpt	images
Deraining	Raindrop	31.87	0.935	MAXIM-2S	14.1M	216G	ckpt	images
Dehazing	RESIDE-Indoor	38.11	0.991	MAXIM-2S	14.1M	216G	ckpt	images
Dehazing	RESIDE-Outdoor	34.19	0.985	MAXIM-2S	14.1M	216G	ckpt	images
Enhancement	LOL	23.43	0.863	MAXIM-2S	14.1M	216G	ckpt	images
Enhancement	FiveK	26.15	0.945	MAXIM-2S	14.1M	216G	ckpt	images

Web Demo

Try the web demo for Image Denoising, Deblurring, Deraining, Dehazing and Enhancement with customised input image here

Demo

First download corresponding checkpoints and then go ahead and run:

Image Denoising (click to expand)

python3 maxim/run_eval.py --task Denoising --ckpt_path ${SIDD_CKPT_PATH} \
  --input_dir maxim/images/Denoising --output_dir maxim/images/Results --has_target=False

Image Deblurring (click to expand)

python3 maxim/run_eval.py --task Deblurring --ckpt_path ${GOPRO_CKPT_PATH} \
  --input_dir maxim/images/Deblurring --output_dir maxim/images/Results --has_target=False

Image Deraining (click to expand)

Rain streak:

python3 maxim/run_eval.py --task Deraining --ckpt_path ${RAIN13K_CKPT_PATH} \
  --input_dir maxim/images/Deraining --output_dir maxim/images/Results --has_target=False

Rain drop:

python3 maxim/run_eval.py --task Deraining --ckpt_path ${RAINDROP_CKPT_PATH} \
  --input_dir maxim/images/Deraining --output_dir maxim/images/Results --has_target=False

Image Dehazing (click to expand)

Indoor:

python3 maxim/run_eval.py --task Dehazing --ckpt_path ${REDISE_INDOOR_CKPT_PATH} \
  --input_dir maxim/images/Dehazing --output_dir maxim/images/Results --has_target=False

Outdoor:

python3 maxim/run_eval.py --task Dehazing --ckpt_path ${REDISE_OUTDOOR_CKPT_PATH} \
  --input_dir maxim/images/Dehazing --output_dir maxim/images/Results --has_target=False

Image Enhancement (click to expand)

Low-light enhancement:

python3 maxim/run_eval.py --task Enhancement --ckpt_path ${LOL_CKPT_PATH} \
  --input_dir maxim/images/Enhancement --output_dir maxim/images/Results --has_target=False

Retouching:

python3 maxim/run_eval.py --task Enhancement --ckpt_path ${FIVEK_CKPT_PATH} \
  --input_dir maxim/images/Enhancement --output_dir maxim/images/Results --has_target=False

Results

Image Denoising (click to expand)

Image Deblurring (click to expand)

Synthetic blur	Realistic blur

Image Deraining (click to expand)

Rain streak	Rain drop

Image Dehazing (click to expand)

Image Enhancement (click to expand)

Citation

Should you find this repository useful, please consider citing:

@article{tu2022maxim,
  title={MAXIM: Multi-Axis MLP for Image Processing},
  author={Tu, Zhengzhong and Talebi, Hossein and Zhang, Han and Yang, Feng and Milanfar, Peyman and Bovik, Alan and Li, Yinxiao},
  journal={CVPR},
  year={2022},
}

Acknowledgement

This repository is built on the vision_transformer and musiq repositories. Our work is also inspired by HiT, MPRNet, and HINet.