Skip to content

Fork of the maxim repository with updated dependencies,and code explanations("MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.)

License

Notifications You must be signed in to change notification settings

theBestPatate/maxim-xplain

 
 

Repository files navigation

PWC PWC PWC PWC PWC PWC

PWC PWC

PWC PWC PWC PWC PWC

PWC PWC

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

Paper Open In Colab slides Poster

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:

Quick Demos

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

Try the Colab here using Open In Colab

Architecture

Model overview

Installation

Changes

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 Project

I 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:

  1. TensorFlow Incompatibility

  2. 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!
  3. 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.
  4. Dependency Updates

    • I added Pillow to the project dependencies to make sure everything functions as it should.

Installation process

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

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

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.

About

Fork of the maxim repository with updated dependencies,and code explanations("MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.)

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 59.1%
  • Jupyter Notebook 40.9%