MAXIM: Multi-Axis MLP for Image Processing (CVPR 2022 Oral, Best Paper Nomination)

This repo hosts the official implementation of the MAXIM models:

"MAXIM: Multi-Axis MLP for Image Processing". CVPR 2022 Oral.
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.

News:

• Jan 8, 2023: Released a pytorch implementation. Check it out here: maxim-pytorch.
• Oct 21, 2022: MAXIM models have been ported to TensorFlow by @sayakpaul. Check it out here: maxim-tf. He also created a couple of Hugging Face Spaces to allow users to quickly try out the different models:
  • Denoising
  • Low-light enhancement
  • Image retouching
  • Dehazing indoors
  • Dehazing outdoors
  • Image deraining
  • Image deblurring
• Sep 8, 2022: our Google AI blog covering both MaxViT and MAXIM is live.
• Apr 25, 2022: Added demos.
  • Colab demo by @deshwalmahesh
  • Replicate web demo .
• Jun 22, 2022: MAXIM selected as 1 of the best paper nomination!
• Mar 29, 2022: MAXIM selected for an oral presentation at CVPR 2022!
• Mar 28, 2022: initial push to Github.
• Mar 3, 2022: paper accepted to CVPR 2022!
• Jan 9, 2022: initial uploads to Arxiv

Quick Demos

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

Try the Colab here using

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 <!--(https://console.cloud.google.com/storage/browser/gresearch/maxim/results/Deblurring/GoPro/)-->
Deblurring	REDS	28.93	0.865	MAXIM-3S	22.2M	339G	ckpt	images <!--(https://console.cloud.google.com/storage/browser/gresearch/maxim/results/Deblurring/REDS/)-->
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

<!-- You can also download most of the training and test datasets we used on [gcloud](https://console.cloud.google.com/storage/browser/gresearch/maxim/datasets/). -->

Demo

First download corresponding checkpoints and then go ahead and run:

<details> <summary><strong>Image Denoising</strong> (click to expand) </summary>

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

</details> <details> <summary><strong>Image Deblurring</strong> (click to expand) </summary>

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

</details> <details> <summary><strong>Image Deraining</strong> (click to expand) </summary>

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

</details> <details> <summary><strong>Image Dehazing</strong> (click to expand) </summary>

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

</details> <details> <summary><strong>Image Enhancement</strong> (click to expand) </summary>

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

</details>

Results

<details> <summary><strong>Image Denoising</strong> (click to expand) </summary> <img src = "https://user-images.githubusercontent.com/43280278/149262475-a73668f2-9fe1-4374-8ed3-4831acca8052.png" width="400"> </details> <details> <summary><strong>Image Deblurring</strong> (click to expand) </summary> <table> <tr> <td> <img src = "https://user-images.githubusercontent.com/43280278/149261823-b77e9513-b3b5-4caf-a0eb-67bf18c2f681.png" width="500"> </td> <td> <img src = "https://user-images.githubusercontent.com/43280278/149261858-24664c33-dc8a-47c3-b84d-ba64b1c05937.png" width="500"> </td> </tr> <tr> <td><p align="center"><b>Synthetic blur</b></p></td> <td><p align="center"><b>Realistic blur</b></p></td> </tr> </table> </details> <details> <summary><strong>Image Deraining</strong> (click to expand) </summary> <table> <tr> <td> <img src = "https://user-images.githubusercontent.com/43280278/149261908-8bce72cf-b343-4bf8-8462-8be363616cfa.png" width="700"> </td> <td> <p align="top"> <img src = "https://user-images.githubusercontent.com/43280278/149262066-7b93538a-2ccc-4ea0-9187-ef1b54734392.png" width="400"> </td> </tr> <tr> <td><p align="center"><b>Rain streak</b></p></td> <td><p align="center"><b>Rain drop</b></p></td> </tr> </table> </details> <details> <summary><strong>Image Dehazing</strong> (click to expand) </summary> <img src = "https://user-images.githubusercontent.com/43280278/149261947-22954827-ce62-44e8-974a-0aa8d94a4bd9.png" width="250"> </details> <details> <summary><strong>Image Enhancement</strong> (click to expand) </summary> <img src = "https://user-images.githubusercontent.com/43280278/149262540-77d16592-9305-4fd7-80c6-b9d30000cc29.png" width="400"> </details>

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.