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:
- 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
.
- Colab demo by @deshwalmahesh
- 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
Try the web demo for Image Denoising, Deblurring, Deraining, Dehazing and Enhancement with customized input image here
The original requirements file was quite permissive and I add to change it to make it work. If you run the original one, you might run into problems with Jax,ml-dtypes,and or Tensorflow. I updated the requirements by restricting the possible versions to make it work and adding Pillow.
More details
# Resolving Compatibility Issues in the ProjectI ran into some compatibility issues while working on the project, and here’s a rundown of the changes I made to get things back on track:
-
TensorFlow Incompatibility
- I discovered some incompatibility issues between TensorFlow and certain versions of NumPy. If you're curious, check out these links for more details:
-
NumPy Compatibility
- I hit an error that I suspect is related to TensorFlow. Here’s the scoop:
- A module compiled with NumPy 1.x just won’t work with NumPy 2.2.5—it can crash and burn.
- To keep things running smoothly with both 1.x and 2.x versions of NumPy, you need to compile your modules with NumPy 2.0.
- The quick fix? Just downgrade NumPy to a 1.x version, and you should be good to go!
- I hit an error that I suspect is related to TensorFlow. Here’s the scoop:
-
Deprecation Notices
- The
float8_e4m3b11
type was deprecated in version 0.2.0 and completely removed in version 0.3.0 of JAX. - Also, the
ml-dtypes
library was an implicit dependency before, but now it’s been made explicit and constrained because one of its deprecated functions was causing some issues in the code.
- The
-
Dependency Updates
- I added Pillow to the project dependencies to make sure everything functions as it should.
Install dependencies:
Using pip
:
pip install -r requirements.txt
Or Using uv
:
uv add -r requirements.txt
Setup project:
the -e option
The -e option is well ... optional but useful if you intend to change the code in this repo. Every changes you will made in the source code will be reflected in the loaded maxim library.Using pip
:
pip install -e .
Or Using uv
:
uv pip install -e
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 |
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
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},
}
This repository is built on the vision_transformer and musiq repositories. Our work is also inspired by HiT, MPRNet, and HINet.