Skip to main content

MAXIM models

· 4 min read

I've released Tensorflow.js ports of the MAXIM family of models, supporting the ability to deblur, denoise, derain, dehaze, retouch, and low-light enhance images.

Motivation​

UpscalerJS was originally created in 2020 with a primary goal of upscaling images. So far, all the models have been exclusively focused on super resolution.

Today I'm releasing a new family of MAXIM models for UpscalerJS that enable a variety of new image enhancement techniques, including:

These models are available in Javascript via UpscalerJS, and run in both the browser and Node.js.

MAXIM​

MAXIM is the architecture at the heart of these new models.

The MAXIM paper (MAXIM: Multi-Axis MLP for Image Processing) was published in 2022, and was nominated as one of the best papers at CVPR 2022. The MAXIM architecture is capable, via training, of supporting a variety of image enhancement tasks. It's also an efficient architecture, making it particularly appealing for JavaScript applications.

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 ... 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.

Google Research published an implementation in Jax, and additional ports were made available in Tensorflow and Pytorch.

Getting Started​

You can run MAXIM models in the browser or Node.js. To get started, install your desired model:

npm install @upscalerjs/maxim-deblurring

And provide the model as an argument to UpscalerJS:

import model from '@upscalerjs/maxim-deblurring'
const upscaler = new Upscaler({
  model,
})

For model-specific instructions, check out the specific model page you're interested in.

Samples​

Below are some samples of each image enhancement task:

Deblurring​

Before

Deblurring Before

After

Deblurring After

Denoising​

Before

Denoising Before

After

Denoising After

Deraining​

Before

Deraining Before

After

Deraining After

Low Light Enhancement​

Before

Low Light Enhancement Before

After

Low Light Enhancement After

Retouching​

Before

Retouching Before

After

Retouching After

Dehazing Indoor​

Before

Dehazing Indoor Before

After

Dehazing Indoor After

Dehazing Outdoor​

Before

Dehazing Outdoor Before

After

Dehazing Outdoor After

Technical Information​

My original attempts at getting this working leveraged both the Jax repo as well as the Tensorflow port. The ported Jax model exhibited close-to-identical fidelity with its Python implementation, but the Tensorflow port was far more performant in the browser, albiet with noticeably inferior fidelity.

Both implementations originally required fixed size inputs in order to port to Tensorflow.js. Fixed size inputs require chunking images, which can break models that rely on a holisitic view of the image, specifically the Dehazing models, Enhancement model, and Retouching model. (For samples of the artifacting this produces and a longer discussion, see this Github issue.)

I modified the MAXIM Jax code to support dynamic image input sizes, and opened a PR here. This PR is also integrated in my fork of the MAXIM code.

(There is also an open PR against the Tensorflow repo exploring dynamic sizes; when it gets merged, I'll explore porting it to Tensorflow.js as well.)

If you'd like to check out the Tensorflow implementation port (noticeably faster, noticeably worse inference, and a fixed size input) these models are available under the maxim-experiments repo in the UpscalerJS repo. Clone the repo, pull the models (dvc pull) and you'll see fixed-input models of 64 and 256 pixel sizes respectively for each task.

Converting the files yourself​

If you'd like to convert the original Jax or Tensorflow model files yourself, instructions are in these Jupyter notebooks. Feel free to open an issue on Github if you run into issues or questions.

The past few years have seen an explosion of innovation in the image enhancement space, and I hope to continue bringing the latest innovations to Javascript. MAXIM is a first step towards enabling Javascript-based image enhancement tasks beyond super resolution in UpscalerJS.

If you have particular models you'd like to see available via UpscalerJS, let me know in Github. If you use MAXIM in your work, I'd love to hear about it! And if you run into questions or find bugs, please open a bug report.