Introduces new Maxim models, including:

- Deblurring
- Deraining
- Denoising
- Low Light Enhancement
- Retouching
- Dehazing (Indoor & Outdoor)

dev/browser/package.json

@@ -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"

dev/browser/specific-model/index.ts

@@ -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`);

docs/blog/2023-08-19-maxim-models.md

@@ -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=).

docs/docs/documentation/troubleshooting.md

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

docs/docusaurus.config.js

@@ -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',
+              },
             ],
           },
         ],

docs/src/components/homepage/header/demo-video/images.ts

@@ -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,
@@ -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'),
@@ -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);

docs/src/components/modelCards/modelCard/modelCard.module.scss

@@ -52,6 +52,8 @@
   img {
     flex: 1;
     width: 50%;
+    aspect-ratio: 1 / 1;
+    object-fit: cover;
   }
 }
 

docs/src/theme/DocItem/Layout/index.js

@@ -65,7 +65,7 @@ export default function DocItemLayout({children}) {
                 {docTOC.mobile}
                 <DocItemContent>
                   {children}
-                  </DocItemContent>
+                </DocItemContent>
                 <DocItemFooter />
               </article>
               <DocItemPaginator />

internals/common/src/models.ts

@@ -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}`);

internals/test-runner/src/ClientsideTestRunner.ts

@@ -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() {

models/default-model/package.json

@@ -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"
     }
@@ -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"

models/default-model/src/cjs.ts

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

models/maxim-deblurring/.gitignore

@@ -0,0 +1,3 @@
+dist
+node_modules
+*.generated.ts

models/maxim-deblurring/.npmignore

@@ -0,0 +1,4 @@
+src
+yarn-error.log
+node_modules
+test