Skip to content

Commit

Permalink
Introduces new Maxim models, including:
Browse files Browse the repository at this point in the history
- Deblurring
- Deraining
- Denoising
- Low Light Enhancement
- Retouching
- Dehazing (Indoor & Outdoor)
  • Loading branch information
thekevinscott committed Nov 21, 2023
1 parent d2e6fcd commit 52daa05
Show file tree
Hide file tree
Showing 410 changed files with 5,101 additions and 57 deletions.
8 changes: 8 additions & 0 deletions dev/browser/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,14 @@
"@upscalerjs/esrgan-slim": "workspace:*",
"@upscalerjs/esrgan-thick": "workspace:*",
"@upscalerjs/pixel-upsampler": "workspace:*",
"@upscalerjs/maxim-experiments": "workspace:*",
"@upscalerjs/maxim-deblurring": "workspace:*",
"@upscalerjs/maxim-deraining": "workspace:*",
"@upscalerjs/maxim-denoising": "workspace:*",
"@upscalerjs/maxim-dehazing-indoor": "workspace:*",
"@upscalerjs/maxim-dehazing-outdoor": "workspace:*",
"@upscalerjs/maxim-enhancement": "workspace:*",
"@upscalerjs/maxim-retouching": "workspace:*",
"seedrandom": "3.0.5",
"upscaler": "workspace:*",
"vite": "^4.4.9"
Expand Down
7 changes: 5 additions & 2 deletions dev/browser/specific-model/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,23 +54,26 @@ const loadModel = async (packageName: string, modelName: string) => {

const img = await makeImg(fixture, `Original: ${packageName}/${modelName}`, 1);
const modelPath = getModelPath(packageName, modelJSON.path);
console.log(modelJSON)
const upscaledImg = await upscaleImage({
...modelJSON,
path: modelPath,
}, img);
}, img, 64, 2);
await makeImg(upscaledImg, `Upscaled: ${packageName}/${modelName}`);
};


const upscaleImage = async (model: ModelDefinition, img: HTMLImageElement | HTMLCanvasElement, patchSize: undefined | number = 64) => {
const upscaleImage = async (model: ModelDefinition, img: HTMLImageElement | HTMLCanvasElement, patchSize?: number, padding?: number) => {
status.innerHTML = 'Starting';
const upscaler = new Upscaler({
model,
});
status.innerHTML = 'Upscaling...';
const start = performance.now();
console.log({ patchSize, padding })
const upscaledImg = await upscaler.upscale(img, {
patchSize,
padding,
progress: console.log,
});
console.log(`Duration: ${((performance.now() - start) / 1000).toFixed(2)}s`);
Expand Down
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified docs/assets/assets/homepage-demo/enhanced/face1.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified docs/assets/assets/homepage-demo/enhanced/face2.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified docs/assets/assets/homepage-demo/enhanced/face3.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified docs/assets/assets/homepage-demo/enhanced/flower.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
152 changes: 152 additions & 0 deletions docs/blog/2023-08-19-maxim-models.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
---
title: MAXIM models
description: New models based on the MAXIM architecture released
slug: maxim-models
authors: kscott
tags: [announcement]
hide_table_of_contents: false
---

I've released Tensorflow.js ports of the [MAXIM family of models](https://arxiv.org/abs/2201.02973), supporting the ability to deblur, denoise, derain, dehaze, retouch, and low-light enhance images.

<!--truncate-->

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

- [Deblurring](https://upscalerjs.com/models/available/maxim-deblurring)
- [Denoising](https://upscalerjs.com/models/available/maxim-denoising)
- [Deraining](https://upscalerjs.com/models/available/maxim-deraining)
- Dehazing (both [indoor](https://upscalerjs.com/models/available/maxim-dehazing-indoor) and [outdoor](https://upscalerjs.com/models/available/maxim-dehazing-outdoor))
- [Low Light Enhancement](https://upscalerjs.com/models/available/maxim-enhancement)
- [Retouching](https://upscalerjs.com/models/available/maxim-retouching)

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](https://arxiv.org/abs/2201.02973)) 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](https://github.com/google-research/maxim), and additional ports were made available in [Tensorflow](https://github.com/sayakpaul/maxim-tf) and [Pytorch](https://github.com/vztu/maxim-pytorch/tree/main/maxim_pytorch).

## Getting Started

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

```bash
npm install @upscalerjs/maxim-deblurring
```

And provide the model as an argument to UpscalerJS:

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

For model-specific instructions, [check out the specific model page](/models) you're interested in.

## Samples

Below are some samples of each image enhancement task:

### Deblurring

*Before*

![Deblurring Before](/assets/sample-images/maxim-deblurring/fixture.png)

*After*

![Deblurring After](/assets/sample-images/maxim-deblurring/result.png)

### Denoising

*Before*

![Denoising Before](/assets/sample-images/maxim-denoising/fixture.png)

*After*

![Denoising After](/assets/sample-images/maxim-denoising/result.png)

### Deraining

*Before*

![Deraining Before](/assets/sample-images/maxim-deraining/fixture.png)

*After*

![Deraining After](/assets/sample-images/maxim-deraining/result.png)

### Low Light Enhancement

*Before*

![Low Light Enhancement Before](/assets/sample-images/maxim-enhancement/fixture.png)

*After*

![Low Light Enhancement After](/assets/sample-images/maxim-enhancement/result.png)

### Retouching

*Before*

![Retouching Before](/assets/sample-images/maxim-retouching/fixture.png)

*After*

![Retouching After](/assets/sample-images/maxim-retouching/result.png)

### Dehazing Indoor

*Before*

![Dehazing Indoor Before](/assets/sample-images/maxim-dehazing-indoor/fixture.png)

*After*

![Dehazing Indoor After](/assets/sample-images/maxim-dehazing-indoor/result.png)

### Dehazing Outdoor

*Before*

![Dehazing Outdoor Before](/assets/sample-images/maxim-dehazing-outdoor/fixture.png)

*After*

![Dehazing Outdoor After](/assets/sample-images/maxim-dehazing-outdoor/result.png)

## Technical Information

My original attempts at getting this working leveraged both the [Jax repo](https://github.com/google-research/maxim/) as well as the [Tensorflow port](https://github.com/sayakpaul/maxim-tf/tree/main). 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](https://github.com/thekevinscott/UpscalerJS/issues/913).)

I modified the MAXIM Jax code to support dynamic image input sizes, [and opened a PR here](https://github.com/google-research/maxim/pull/41). [This PR is also integrated in my fork of the MAXIM code](https://github.com/thekevinscott/maxim).

(There is also an [open PR against the Tensorflow repo exploring dynamic sizes](https://github.com/sayakpaul/maxim-tf/pull/24); 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](https://github.com/upscalerjs/maxim). 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](https://github.com/thekevinscott/UpscalerJS/discussions/new?category=ideas). If you use MAXIM in your work, [I'd love to hear about it](https://github.com/thekevinscott/UpscalerJS/discussions/new?category=show-and-tell)! And if you run into questions or find bugs, [please open a bug report](https://github.com/thekevinscott/UpscalerJS/issues/new?assignees=thekevinscott&labels=&projects=&template=bug_report.md&title=).
12 changes: 5 additions & 7 deletions docs/docs/documentation/troubleshooting.md
Original file line number Diff line number Diff line change
Expand Up @@ -65,27 +65,25 @@ This likely means one of two things:
- You are using `upscaler`, instead of `upscaler/node`; [check out the guide on Node.js here](/documentation/guides/node/nodejs).
- You are using `import` syntax instead of `require` syntax; if so, try switching to `require('upscaler')`. For more information on this, [see this Github issue](https://github.com/thekevinscott/UpscalerJS/issues/554#issuecomment-1344108954).

## Missing Model Scale
## Missing Model Path

If you see an error like:

```
Error: You must provide a "scale" when providing a model definition
Error: You must provide a "path" when providing a model definition
```

You've passed a `null` or `undefined` scale argument in the `model` argument to UpscalerJS:
You've passed a `null` or `undefined` path argument in the `model` argument to UpscalerJS:

```javascript
const upscaler = new Upscaler({
model: {
scale: null,
path: null,
},
})
```

Every model must have an explicit `scale` defined.

Ensure you pass a valid `scale` argument in the `model` payload. [See the guide on models for more information](/documentation/guides/browser/models).
Ensure you pass a valid `path` argument in the `model` payload. [See the guide on models for more information](/documentation/guides/browser/models).

## Invalid Warmup Value

Expand Down
24 changes: 24 additions & 0 deletions docs/docusaurus.config.js
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,30 @@ const config = {
label: 'Super Resolution',
to: '/models#upscaling',
},
{
label: 'Deblurring',
to: '/models#deblurring',
},
{
label: 'Denoising',
to: '/models#denoising',
},
{
label: 'Deraining',
to: '/models#deraining',
},
{
label: 'Dehazing',
to: '/models#dehazing',
},
{
label: 'Low Light Enhancement',
to: '/models#low-light-enhancement',
},
{
label: 'Retouching',
to: '/models#retouching',
},
],
},
],
Expand Down
27 changes: 27 additions & 0 deletions docs/src/components/homepage/header/demo-video/images.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,26 @@
import { IExampleImage } from "./example-images/example-images";

const makeMAXIM = (model: string, action: string): IExampleImage => ({
sizes: {
original: 256,
enhanced: 256,
},
original: {
src: `/assets/homepage-demo/originals/${model}.png`,
labels: {
short: 'Original',
long: 'Original',
},
},
enhanced: {
src: `/assets/homepage-demo/enhanced/${model}.png`,
labels: {
short: `@upscalerjs/maxim-${model}`,
long: `${action} using [\`@upscalerjs/maxim-${model}\`](/models/available/maxim-${model}) model`,
},
},
});

const makeESRGAN = (image: string): IExampleImage => ({
sizes: {
original: 128,
Expand All @@ -22,6 +43,9 @@ const makeESRGAN = (image: string): IExampleImage => ({
});

const images: Record<string, IExampleImage> = {
maximDenoising: makeMAXIM('denoising', 'Denoised'),
maximEnhancement: makeMAXIM('enhancement', 'Enhanced'),
maximDeblurring: makeMAXIM('deblurring', 'Deblurred'),
esrganFlower: makeESRGAN('flower'),
esrganFace3: makeESRGAN('face3'),
esrganFace2: makeESRGAN('face2'),
Expand All @@ -31,6 +55,9 @@ const images: Record<string, IExampleImage> = {
export const IMAGES: IExampleImage[] = [
images.esrganFlower,
images.esrganFace2,
images.maximEnhancement,
images.esrganFace3,
images.maximDeblurring,
images.esrganFace1,
images.maximDenoising,
].filter(Boolean);
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@
img {
flex: 1;
width: 50%;
aspect-ratio: 1 / 1;
object-fit: cover;
}
}

Expand Down
2 changes: 1 addition & 1 deletion docs/src/theme/DocItem/Layout/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ export default function DocItemLayout({children}) {
{docTOC.mobile}
<DocItemContent>
{children}
</DocItemContent>
</DocItemContent>
<DocItemFooter />
</article>
<DocItemPaginator />
Expand Down
9 changes: 7 additions & 2 deletions internals/common/src/models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,11 +56,16 @@ export const getUMDNames = async (packageName: string): Promise<Record<string, s
}
}

const getAllAvailableModels = async (packageName: string): Promise<AvailableModel[]> => {
const getAllAvailableModels = async (packageName: string, includeIndex = false): Promise<AvailableModel[]> => {
const modelPackageDir = path.resolve(MODELS_DIR, packageName);
const umdNames = await getUMDNames(modelPackageDir);
const packageJSONExports = await getPackageJSONExports(modelPackageDir);
return packageJSONExports.filter(k => k[0] !== '.').map(([key, value]) => {
return packageJSONExports.filter(k => {
if (packageJSONExports.length > 1) {
return k[0] !== '.';
}
return true;
}).map(([key, value]) => {
const umdName = umdNames[key];
if (umdName === undefined) {
throw new Error(`No UMD name defined for ${packageName}/umd-names.json for ${key}`);
Expand Down
7 changes: 7 additions & 0 deletions internals/test-runner/src/ClientsideTestRunner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -219,9 +219,16 @@ export class ClientsideTestRunner {
}

public async startBrowser() {
// launch handles launching an instance and then connecting to it
this.browser = await launch({
headless: 'new',
protocolTimeout: 180_000 * 2,
});

// connect is for connecting to an already running instance
// this.browser = await connect({

// });
}

private _attachLogger() {
Expand Down
6 changes: 5 additions & 1 deletion models/default-model/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
},
"exports": {
".": {
"umd": "./dist/umd/index.min.js",
"require": "./dist/cjs/models/default-model/src/cjs.js",
"import": "./dist/esm/models/default-model/src/index.js"
}
Expand Down Expand Up @@ -63,7 +64,10 @@
},
"wireit": {
"lint": {
"command": "eslint -c ../.eslintrc.js src --ext .ts"
"command": "eslint -c ../.eslintrc.js src --ext .ts",
"dependencies": [
"scaffold"
]
},
"prepublishOnly": {
"command": "pnpm lint && pnpm build && pnpm validate:build"
Expand Down
4 changes: 2 additions & 2 deletions models/default-model/src/cjs.ts
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
import { default as DefaultModel, } from './index';
module.exports = DefaultModel; // eslint-disable-line
import { default as model, } from './index';
module.exports = model; // eslint-disable-line
3 changes: 3 additions & 0 deletions models/maxim-deblurring/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
dist
node_modules
*.generated.ts
4 changes: 4 additions & 0 deletions models/maxim-deblurring/.npmignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
src
yarn-error.log
node_modules
test
Loading

0 comments on commit 52daa05

Please sign in to comment.