108 |
115 |
116 | Image Denoising (click to expand)
109 | 110 | ``` 111 | python3 maxim/run_eval.py --task Denoising --ckpt_path ${SIDD_CKPT_PATH} \ 112 | --input_dir maxim/images/Denoising --output_dir maxim/images/Results --has_target=False 113 | ``` 114 |
117 |
124 |
125 | Image Deblurring (click to expand)
118 | 119 | ``` 120 | python3 maxim/run_eval.py --task Deblurring --ckpt_path ${GOPRO_CKPT_PATH} \ 121 | --input_dir maxim/images/Deblurring --output_dir maxim/images/Results --has_target=False 122 | ``` 123 |
126 |
140 |
141 | Image Deraining (click to expand)
127 | 128 | Rain streak: 129 | ``` 130 | python3 maxim/run_eval.py --task Deraining --ckpt_path ${RAIN13K_CKPT_PATH} \ 131 | --input_dir maxim/images/Deraining --output_dir maxim/images/Results --has_target=False 132 | ``` 133 | 134 | Rain drop: 135 | ``` 136 | python3 maxim/run_eval.py --task Deraining --ckpt_path ${RAINDROP_CKPT_PATH} \ 137 | --input_dir maxim/images/Deraining --output_dir maxim/images/Results --has_target=False 138 | ``` 139 |
142 |
156 |
157 | Image Dehazing (click to expand)
143 | 144 | Indoor: 145 | ``` 146 | python3 maxim/run_eval.py --task Dehazing --ckpt_path ${REDISE_INDOOR_CKPT_PATH} \ 147 | --input_dir maxim/images/Dehazing --output_dir maxim/images/Results --has_target=False 148 | ``` 149 | 150 | Outdoor: 151 | ``` 152 | python3 maxim/run_eval.py --task Dehazing --ckpt_path ${REDISE_OUTDOOR_CKPT_PATH} \ 153 | --input_dir maxim/images/Dehazing --output_dir maxim/images/Results --has_target=False 154 | ``` 155 |
158 |
172 |
173 | ## Results
174 |
175 | Image Enhancement (click to expand)
159 | 160 | Low-light enhancement: 161 | ``` 162 | python3 maxim/run_eval.py --task Enhancement --ckpt_path ${LOL_CKPT_PATH} \ 163 | --input_dir maxim/images/Enhancement --output_dir maxim/images/Results --has_target=False 164 | ``` 165 | 166 | Retouching: 167 | ``` 168 | python3 maxim/run_eval.py --task Enhancement --ckpt_path ${FIVEK_CKPT_PATH} \ 169 | --input_dir maxim/images/Enhancement --output_dir maxim/images/Results --has_target=False 170 | ``` 171 |
176 | 
179 |
180 |
181 | Image Denoising (click to expand)
177 | 178 |
179 |
182 |
185 |
194 |
195 |
196 | Image Deblurring (click to expand)
183 | 184 |  |
187 |  |
188 |
Synthetic blur |
191 | Realistic blur |
192 |
197 |
200 |
209 |
210 |
211 | Image Deraining (click to expand)
198 | 199 |  |
202 | |
203 |
Rain streak |
206 | Rain drop |
207 |
212 | 
215 |
216 |
217 | Image Dehazing (click to expand)
213 | 214 |
215 |
218 | 
221 |
222 |
223 | ## Citation
224 | Should you find this repository useful, please consider citing:
225 | ```
226 | @article{tu2022maxim,
227 | title={MAXIM: Multi-Axis MLP for Image Processing},
228 | author={Tu, Zhengzhong and Talebi, Hossein and Zhang, Han and Yang, Feng and Milanfar, Peyman and Bovik, Alan and Li, Yinxiao},
229 | journal={CVPR},
230 | year={2022},
231 | }
232 | ```
233 |
234 | ## Acknowledgement
235 |
236 | This repository is built on the [vision_transformer](https://github.com/google-research/vision_transformer) and [musiq](https://github.com/google-research/google-research/tree/master/musiq) repositories. Our work is also inspired by [HiT](https://github.com/google-research/hit-gan), [MPRNet](https://github.com/swz30/MPRNet), and [HINet](https://github.com/megvii-model/HINet).
237 |
--------------------------------------------------------------------------------
/cog.yaml:
--------------------------------------------------------------------------------
1 | build:
2 | cuda: "11.2"
3 | gpu: true
4 | python_version: "3.8"
5 | system_packages:
6 | - "libgl1-mesa-glx"
7 | - "libglib2.0-0"
8 | python_packages:
9 | - "numpy==1.21.1"
10 | - "ipython==7.21.0"
11 | - "absl-py==1.0.0"
12 | - "chex==0.1.3"
13 | - "clu==0.0.6"
14 | - "einops==0.4.1"
15 | - "flax==0.4.1"
16 | - "ml-collections==0.1.1"
17 | - "pandas==1.4.2"
18 | - "tensorflow==2.8.0"
19 | run:
20 | - pip install --upgrade pip
21 | - pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
22 |
23 | predict: "maxim/predict.py:Predictor"
--------------------------------------------------------------------------------
/colab_inference_demo.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "markdown",
5 | "metadata": {
6 | "id": "cSLfR2K1fRgR"
7 | },
8 | "source": [
9 | "# Inference Notebook\n",
10 | "\n",
11 | "[MAXIM: Multi-Axis MLP for Image Processing (CVPR 2022 Oral)](https://github.com/google-research/maxim)\n",
12 | "\n",
13 | "**This is just the inference code. Maximum you can do is to come in with your images and get results using trained models**"
14 | ]
15 | },
16 | {
17 | "cell_type": "markdown",
18 | "metadata": {
19 | "id": "vXxVwI1SfJc-"
20 | },
21 | "source": [
22 | "# Clone repo and install dependencies"
23 | ]
24 | },
25 | {
26 | "cell_type": "code",
27 | "execution_count": null,
28 | "metadata": {
29 | "colab": {
30 | "base_uri": "https://localhost:8080/"
31 | },
32 | "id": "RbskPSxADHXM",
33 | "outputId": "35c15104-7636-4ef6-cfda-f8ffe0e7e3fe"
34 | },
35 | "outputs": [
36 | {
37 | "output_type": "stream",
38 | "name": "stdout",
39 | "text": [
40 | ""
41 | ]
42 | }
43 | ],
44 | "source": [
45 | "! git clone https://github.com/google-research/maxim/\n",
46 | "%cd ./maxim\n",
47 | "\n",
48 | "!pip install -r requirements.txt\n",
49 | "!pip install --upgrade jax\n",
50 | "! pip install gdown\n",
51 | "\n",
52 | "!python setup.py build\n",
53 | "! python setup.py install\n",
54 | "\n",
55 | "# https://console.cloud.google.com/storage/browser/gresearch/maxim/ckpt/Enhancement/FiveK;tab=objects?prefix=&forceOnObjectsSortingFiltering=false"
56 | ]
57 | },
58 | {
59 | "cell_type": "markdown",
60 | "metadata": {
61 | "id": "l06nrsoVdFRA"
62 | },
63 | "source": [
64 | "# Imports and Defaults\n",
65 | "Imports from libraries and from the modules written by authors of the repo\n",
66 | "\n"
67 | ]
68 | },
69 | {
70 | "cell_type": "code",
71 | "execution_count": null,
72 | "metadata": {
73 | "id": "LEbaNP5IdNOQ"
74 | },
75 | "outputs": [],
76 | "source": [
77 | "from google.colab import drive # works only for colab\n",
78 | "from PIL import Image\n",
79 | "\n",
80 | "import matplotlib.pyplot as plt\n",
81 | "import collections\n",
82 | "import importlib\n",
83 | "import io\n",
84 | "import os\n",
85 | "import math\n",
86 | "import requests\n",
87 | "from tqdm import tqdm\n",
88 | "import gdown # to download weights from Drive\n",
89 | "\n",
90 | "import flax\n",
91 | "import jax.numpy as jnp\n",
92 | "import ml_collections\n",
93 | "import numpy as np\n",
94 | "import tensorflow as tf\n",
95 | "from jax.experimental import jax2tf\n",
96 | "\n",
97 | "\n",
98 | "# below code lines are from run_eval.py\n",
99 | "_MODEL_FILENAME = 'maxim'\n",
100 | "\n",
101 | "_MODEL_VARIANT_DICT = {\n",
102 | " 'Denoising': 'S-3',\n",
103 | " 'Deblurring': 'S-3',\n",
104 | " 'Deraining': 'S-2',\n",
105 | " 'Dehazing': 'S-2',\n",
106 | " 'Enhancement': 'S-2',\n",
107 | "}\n",
108 | "\n",
109 | "_MODEL_CONFIGS = {\n",
110 | " 'variant': '',\n",
111 | " 'dropout_rate': 0.0,\n",
112 | " 'num_outputs': 3,\n",
113 | " 'use_bias': True,\n",
114 | " 'num_supervision_scales': 3,\n",
115 | "}\n"
116 | ]
117 | },
118 | {
119 | "cell_type": "markdown",
120 | "metadata": {
121 | "id": "EILubkkjc1P5"
122 | },
123 | "source": [
124 | "# Link Google Drive for data input and output \n",
125 | "Not necessary but ease of use for Data input / Output"
126 | ]
127 | },
128 | {
129 | "cell_type": "code",
130 | "execution_count": null,
131 | "metadata": {
132 | "id": "w4qz8CMTJonN"
133 | },
134 | "outputs": [],
135 | "source": [
136 | "# drive.mount('/content/gdrive/',)"
137 | ]
138 | },
139 | {
140 | "cell_type": "markdown",
141 | "metadata": {
142 | "id": "54_eNSRJdUHz"
143 | },
144 | "source": [
145 | "# Helpers"
146 | ]
147 | },
148 | {
149 | "cell_type": "code",
150 | "execution_count": null,
151 | "metadata": {
152 | "id": "36QlK8Rfk5ai"
153 | },
154 | "outputs": [],
155 | "source": [
156 | "def sizeof_fmt(size, suffix='B'):\n",
157 | " \"\"\"Get human readable file size.\n",
158 | " Args:\n",
159 | " size (int): File size.\n",
160 | " suffix (str): Suffix. Default: 'B'.\n",
161 | " Return:\n",
162 | " str: Formated file siz.\n",
163 | " \"\"\"\n",
164 | " for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:\n",
165 | " if abs(size) < 1024.0:\n",
166 | " return f'{size:3.1f} {unit}{suffix}'\n",
167 | " size /= 1024.0\n",
168 | " return f'{size:3.1f} Y{suffix}'\n",
169 | "\n",
170 | "\n",
171 | "def download_file_from_google_drive(file_id, save_path):\n",
172 | " \"\"\"Download files from google drive.\n",
173 | "\n",
174 | " Ref:\n",
175 | " https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive # noqa E501\n",
176 | "\n",
177 | " Args:\n",
178 | " file_id (str): File id.\n",
179 | " save_path (str): Save path.\n",
180 | " \"\"\"\n",
181 | "\n",
182 | " session = requests.Session()\n",
183 | " URL = 'https://docs.google.com/uc?export=download'\n",
184 | " params = {'id': file_id}\n",
185 | "\n",
186 | " response = session.get(URL, params=params, stream=True)\n",
187 | " token = get_confirm_token(response)\n",
188 | " if token:\n",
189 | " params['confirm'] = token\n",
190 | " response = session.get(URL, params=params, stream=True)\n",
191 | "\n",
192 | " # get file size\n",
193 | " response_file_size = session.get(\n",
194 | " URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})\n",
195 | " if 'Content-Range' in response_file_size.headers:\n",
196 | " file_size = int(\n",
197 | " response_file_size.headers['Content-Range'].split('/')[1])\n",
198 | " else:\n",
199 | " file_size = None\n",
200 | "\n",
201 | " save_response_content(response, save_path, file_size)\n",
202 | "\n",
203 | "\n",
204 | "def get_confirm_token(response):\n",
205 | " for key, value in response.cookies.items():\n",
206 | " if key.startswith('download_warning'):\n",
207 | " return value\n",
208 | " return None\n",
209 | "\n",
210 | "\n",
211 | "def save_response_content(response,\n",
212 | " destination,\n",
213 | " file_size=None,\n",
214 | " chunk_size=32768):\n",
215 | " if file_size is not None:\n",
216 | " pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')\n",
217 | "\n",
218 | " readable_file_size = sizeof_fmt(file_size)\n",
219 | " else:\n",
220 | " pbar = None\n",
221 | "\n",
222 | " with open(destination, 'wb') as f:\n",
223 | " downloaded_size = 0\n",
224 | " for chunk in response.iter_content(chunk_size):\n",
225 | " downloaded_size += chunk_size\n",
226 | " if pbar is not None:\n",
227 | " pbar.update(1)\n",
228 | " pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} '\n",
229 | " f'/ {readable_file_size}')\n",
230 | " if chunk: # filter out keep-alive new chunks\n",
231 | " f.write(chunk)\n",
232 | " if pbar is not None:\n",
233 | " pbar.close()"
234 | ]
235 | },
236 | {
237 | "cell_type": "code",
238 | "execution_count": null,
239 | "metadata": {
240 | "id": "TY2GyIbn6j85"
241 | },
242 | "outputs": [],
243 | "source": [
244 | "\n",
245 | "def resize(path, new_width_height = 1280, save_image = False, convert_RGB = True, clip_full_hd = False, quality = 100):\n",
246 | " '''\n",
247 | " Resize and return Given Image\n",
248 | " args:\n",
249 | " path: Image Path\n",
250 | " new_width_height = Reshaped image's width and height. # If integer is given, it'll keep the aspect ratio as it is by shrinking the Bigger dimension (width or height) to the max of new_width_height and then shring the smaller dimension accordingly \n",
251 | " save_image = Whether to save the image or not\n",
252 | " convert_RGB: Whether to Convert the RGBA image to RGB (by default backgroud is white)\n",
253 | " '''\n",
254 | " image = Image.open(path)\n",
255 | " w, h = image.size\n",
256 | "\n",
257 | " fixed_size = new_width_height if isinstance(new_width_height, int) else False\n",
258 | "\n",
259 | " if fixed_size:\n",
260 | " if h > w:\n",
261 | " fixed_height = fixed_size\n",
262 | " height_percent = (fixed_height / float(h))\n",
263 | " width_size = int((float(w) * float(height_percent)))\n",
264 | " image = image.resize((width_size, fixed_height), Image.NEAREST)\n",
265 | "\n",
266 | " else:\n",
267 | " fixed_width = fixed_size\n",
268 | " width_percent = (fixed_width / float(w))\n",
269 | " height_size = int((float(h) * float(width_percent)))\n",
270 | " image = image.resize((fixed_width, height_size), Image.NEAREST) # Try Image.ANTIALIAS inplace of Image.NEAREST\n",
271 | "\n",
272 | " else:\n",
273 | " image = image.resize(new_width_height)\n",
274 | "\n",
275 | " if image.mode == \"RGBA\" and convert_RGB:\n",
276 | " # image.load() # required for png.split()\n",
277 | " # new = Image.new(\"RGB\", image.size, (255, 255, 255)) # White Background\n",
278 | " # image = new.paste(image, mask=image.split()[3]) # 3 is the alpha channel\n",
279 | "\n",
280 | " new = Image.new(\"RGBA\", image.size, \"WHITE\") # Create a white rgba background\n",
281 | " new.paste(image, (0, 0), image) # Paste the image on the background.\n",
282 | " image = new.convert('RGB')\n",
283 | "\n",
284 | " if save_image:\n",
285 | " image.save(path, quality = quality)\n",
286 | "\n",
287 | " return image\n",
288 | "\n",
289 | "\n",
290 | "class DummyFlags():\n",
291 | " def __init__(self, ckpt_path:str, task:str, input_dir: str = \"./maxim/images/Enhancement\", output_dir:str = \"./maxim/images/Results\", has_target:bool = False, save_images:bool = True, geometric_ensemble:bool = False):\n",
292 | " '''\n",
293 | " Builds the dummy flags which replicates the behaviour of Terminal CLI execution (same as ArgParse)\n",
294 | " args:\n",
295 | " ckpt_path: Saved Model CheckPoint: Find all the checkpoints for pre trained models at https://console.cloud.google.com/storage/browser/gresearch/maxim/ckpt/\n",
296 | " task: Task for which the model waas trained. Each task uses different Data and Checkpoints. Find the details of tasks and respective checkpoints details at: https://github.com/google-research/maxim#results-and-pre-trained-models\n",
297 | " input_dir: Input Directory. We do not need it here as we are directly passing one image at a time\n",
298 | " output_dir: Also not needed in out code\n",
299 | " has_target: Used to calculate PSNR and SSIM calculation. Not needed in our case\n",
300 | " save_images: Used in CLI command where images were saved in loop. Not needed in our case\n",
301 | " geometric_ensemble: Was used in training part and as it is just an Inference part, it is not needed\n",
302 | "\n",
303 | " '''\n",
304 | " self.ckpt_path = ckpt_path\n",
305 | " self.task = task\n",
306 | " self.input_dir = input_dir\n",
307 | " self.output_dir = output_dir\n",
308 | " self.has_target = has_target\n",
309 | " self.save_images = save_images\n",
310 | " self.geometric_ensemble = geometric_ensemble\n"
311 | ]
312 | },
313 | {
314 | "cell_type": "markdown",
315 | "metadata": {
316 | "id": "JkgUJDR0daUP"
317 | },
318 | "source": [
319 | "# Refactored code from authors (`run_eval.py`)\n",
320 | "\n",
321 | "**NOTE**: This is not my code. I just changed the structure, redirected dependencies within modules, removed redundant imports and code and bla bla bla...."
322 | ]
323 | },
324 | {
325 | "cell_type": "code",
326 | "execution_count": null,
327 | "metadata": {
328 | "id": "9oUVYnXQK_WV"
329 | },
330 | "outputs": [],
331 | "source": [
332 | "# Copyright 2022 Google LLC.\n",
333 | "#\n",
334 | "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
335 | "# you may not use this file except in compliance with the License.\n",
336 | "# You may obtain a copy of the License at\n",
337 | "#\n",
338 | "# http://www.apache.org/licenses/LICENSE-2.0\n",
339 | "#\n",
340 | "# Unless required by applicable law or agreed to in writing, software\n",
341 | "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
342 | "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
343 | "# See the License for the specific language governing permissions and\n",
344 | "# limitations under the License.\n",
345 | "\n",
346 | "\n",
347 | "def recover_tree(keys, values):\n",
348 | " \"\"\"Recovers a tree as a nested dict from flat names and values.\n",
349 | "\n",
350 | " This function is useful to analyze checkpoints that are saved by our programs\n",
351 | " without need to access the exact source code of the experiment. In particular,\n",
352 | " it can be used to extract an reuse various subtrees of the scheckpoint, e.g.\n",
353 | " subtree of parameters.\n",
354 | " Args:\n",
355 | " keys: a list of keys, where '/' is used as separator between nodes.\n",
356 | " values: a list of leaf values.\n",
357 | " Returns:\n",
358 | " A nested tree-like dict.\n",
359 | " \"\"\"\n",
360 | " tree = {}\n",
361 | " sub_trees = collections.defaultdict(list)\n",
362 | " for k, v in zip(keys, values):\n",
363 | " if '/' not in k:\n",
364 | " tree[k] = v\n",
365 | " else:\n",
366 | " k_left, k_right = k.split('/', 1)\n",
367 | " sub_trees[k_left].append((k_right, v))\n",
368 | " for k, kv_pairs in sub_trees.items():\n",
369 | " k_subtree, v_subtree = zip(*kv_pairs)\n",
370 | " tree[k] = recover_tree(k_subtree, v_subtree)\n",
371 | " return tree\n",
372 | "\n",
373 | "\n",
374 | "def mod_padding_symmetric(image, factor=64):\n",
375 | " \"\"\"Padding the image to be divided by factor.\"\"\"\n",
376 | " height, width = image.shape[0], image.shape[1]\n",
377 | " height_pad, width_pad = ((height + factor) // factor) * factor, (\n",
378 | " (width + factor) // factor) * factor\n",
379 | " padh = height_pad - height if height % factor != 0 else 0\n",
380 | " padw = width_pad - width if width % factor != 0 else 0\n",
381 | " image = jnp.pad(\n",
382 | " image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)],\n",
383 | " mode='reflect')\n",
384 | " return image\n",
385 | "\n",
386 | "\n",
387 | "def get_params(ckpt_path):\n",
388 | " \"\"\"Get params checkpoint.\"\"\"\n",
389 | "\n",
390 | " with tf.io.gfile.GFile(ckpt_path, 'rb') as f:\n",
391 | " data = f.read()\n",
392 | " values = np.load(io.BytesIO(data))\n",
393 | " params = recover_tree(*zip(*values.items()))\n",
394 | " params = params['opt']['target']\n",
395 | "\n",
396 | " return params\n",
397 | "\n",
398 | "\n",
399 | "def calculate_psnr(img1, img2, crop_border, test_y_channel=False):\n",
400 | " \"\"\"Calculate PSNR (Peak Signal-to-Noise Ratio).\n",
401 | "\n",
402 | " Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio\n",
403 | " Args:\n",
404 | " img1 (ndarray): Images with range [0, 255].\n",
405 | " img2 (ndarray): Images with range [0, 255].\n",
406 | " crop_border (int): Cropped pixels in each edge of an image. These\n",
407 | " pixels are not involved in the PSNR calculation.\n",
408 | " test_y_channel (bool): Test on Y channel of YCbCr. Default: False.\n",
409 | " Returns:\n",
410 | " float: psnr result.\n",
411 | " \"\"\"\n",
412 | " assert img1.shape == img2.shape, (\n",
413 | " f'Image shapes are differnet: {img1.shape}, {img2.shape}.')\n",
414 | " img1 = img1.astype(np.float64)\n",
415 | " img2 = img2.astype(np.float64)\n",
416 | "\n",
417 | " if crop_border != 0:\n",
418 | " img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]\n",
419 | " img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]\n",
420 | "\n",
421 | " if test_y_channel:\n",
422 | " img1 = to_y_channel(img1)\n",
423 | " img2 = to_y_channel(img2)\n",
424 | "\n",
425 | " mse = np.mean((img1 - img2)**2)\n",
426 | " if mse == 0:\n",
427 | " return float('inf')\n",
428 | " return 20. * np.log10(255. / np.sqrt(mse))\n",
429 | "\n",
430 | "\n",
431 | "def _convert_input_type_range(img):\n",
432 | " \"\"\"Convert the type and range of the input image.\n",
433 | "\n",
434 | " It converts the input image to np.float32 type and range of [0, 1].\n",
435 | " It is mainly used for pre-processing the input image in colorspace\n",
436 | " convertion functions such as rgb2ycbcr and ycbcr2rgb.\n",
437 | " Args:\n",
438 | " img (ndarray): The input image. It accepts:\n",
439 | " 1. np.uint8 type with range [0, 255];\n",
440 | " 2. np.float32 type with range [0, 1].\n",
441 | " Returns:\n",
442 | " (ndarray): The converted image with type of np.float32 and range of\n",
443 | " [0, 1].\n",
444 | " \"\"\"\n",
445 | " img_type = img.dtype\n",
446 | " img = img.astype(np.float32)\n",
447 | " if img_type == np.float32:\n",
448 | " pass\n",
449 | " elif img_type == np.uint8:\n",
450 | " img /= 255.\n",
451 | " else:\n",
452 | " raise TypeError('The img type should be np.float32 or np.uint8, '\n",
453 | " f'but got {img_type}')\n",
454 | " return img\n",
455 | "\n",
456 | "\n",
457 | "def _convert_output_type_range(img, dst_type):\n",
458 | " \"\"\"Convert the type and range of the image according to dst_type.\n",
459 | "\n",
460 | " It converts the image to desired type and range. If `dst_type` is np.uint8,\n",
461 | " images will be converted to np.uint8 type with range [0, 255]. If\n",
462 | " `dst_type` is np.float32, it converts the image to np.float32 type with\n",
463 | " range [0, 1].\n",
464 | " It is mainly used for post-processing images in colorspace convertion\n",
465 | " functions such as rgb2ycbcr and ycbcr2rgb.\n",
466 | " Args:\n",
467 | " img (ndarray): The image to be converted with np.float32 type and\n",
468 | " range [0, 255].\n",
469 | " dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it\n",
470 | " converts the image to np.uint8 type with range [0, 255]. If\n",
471 | " dst_type is np.float32, it converts the image to np.float32 type\n",
472 | " with range [0, 1].\n",
473 | " Returns:\n",
474 | " (ndarray): The converted image with desired type and range.\n",
475 | " \"\"\"\n",
476 | " if dst_type not in (np.uint8, np.float32):\n",
477 | " raise TypeError('The dst_type should be np.float32 or np.uint8, '\n",
478 | " f'but got {dst_type}')\n",
479 | " if dst_type == np.uint8:\n",
480 | " img = img.round()\n",
481 | " else:\n",
482 | " img /= 255.\n",
483 | "\n",
484 | " return img.astype(dst_type)\n",
485 | "\n",
486 | "\n",
487 | "def rgb2ycbcr(img, y_only=False):\n",
488 | " \"\"\"Convert a RGB image to YCbCr image.\n",
489 | "\n",
490 | " This function produces the same results as Matlab's `rgb2ycbcr` function.\n",
491 | " It implements the ITU-R BT.601 conversion for standard-definition\n",
492 | " television. See more details in\n",
493 | " https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.\n",
494 | " It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.\n",
495 | " In OpenCV, it implements a JPEG conversion. See more details in\n",
496 | " https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.\n",
497 | "\n",
498 | " Args:\n",
499 | " img (ndarray): The input image. It accepts:\n",
500 | " 1. np.uint8 type with range [0, 255];\n",
501 | " 2. np.float32 type with range [0, 1].\n",
502 | " y_only (bool): Whether to only return Y channel. Default: False.\n",
503 | " Returns:\n",
504 | " ndarray: The converted YCbCr image. The output image has the same type\n",
505 | " and range as input image.\n",
506 | " \"\"\"\n",
507 | " img_type = img.dtype\n",
508 | " img = _convert_input_type_range(img)\n",
509 | " if y_only:\n",
510 | " out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0\n",
511 | " else:\n",
512 | " out_img = np.matmul(img,\n",
513 | " [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],\n",
514 | " [24.966, 112.0, -18.214]]) + [16, 128, 128]\n",
515 | " out_img = _convert_output_type_range(out_img, img_type)\n",
516 | " return out_img\n",
517 | "\n",
518 | "\n",
519 | "def to_y_channel(img):\n",
520 | " \"\"\"Change to Y channel of YCbCr.\n",
521 | "\n",
522 | " Args:\n",
523 | " img (ndarray): Images with range [0, 255].\n",
524 | " Returns:\n",
525 | " (ndarray): Images with range [0, 255] (float type) without round.\n",
526 | " \"\"\"\n",
527 | " img = img.astype(np.float32) / 255.\n",
528 | " if img.ndim == 3 and img.shape[2] == 3:\n",
529 | " img = rgb2ycbcr(img, y_only=True)\n",
530 | " img = img[..., None]\n",
531 | " return img * 255.\n",
532 | "\n",
533 | "\n",
534 | "def augment_image(image, times=8):\n",
535 | " \"\"\"Geometric augmentation.\"\"\"\n",
536 | " if times == 4: # only rotate image\n",
537 | " images = []\n",
538 | " for k in range(0, 4):\n",
539 | " images.append(np.rot90(image, k=k))\n",
540 | " images = np.stack(images, axis=0)\n",
541 | " elif times == 8: # roate and flip image\n",
542 | " images = []\n",
543 | " for k in range(0, 4):\n",
544 | " images.append(np.rot90(image, k=k))\n",
545 | " image = np.fliplr(image)\n",
546 | " for k in range(0, 4):\n",
547 | " images.append(np.rot90(image, k=k))\n",
548 | " images = np.stack(images, axis=0)\n",
549 | " else:\n",
550 | " raise Exception(f'Error times: {times}')\n",
551 | " return images\n",
552 | "\n",
553 | "\n",
554 | "def deaugment_image(images, times=8):\n",
555 | " \"\"\"Reverse the geometric augmentation.\"\"\"\n",
556 | "\n",
557 | " if times == 4: # only rotate image\n",
558 | " image = []\n",
559 | " for k in range(0, 4):\n",
560 | " image.append(np.rot90(images[k], k=4-k))\n",
561 | " image = np.stack(image, axis=0)\n",
562 | " image = np.mean(image, axis=0)\n",
563 | " elif times == 8: # roate and flip image\n",
564 | " image = []\n",
565 | " for k in range(0, 4):\n",
566 | " image.append(np.rot90(images[k], k=4-k))\n",
567 | " for k in range(0, 4):\n",
568 | " image.append(np.fliplr(np.rot90(images[4+k], k=4-k)))\n",
569 | " image = np.mean(image, axis=0)\n",
570 | " else:\n",
571 | " raise Exception(f'Error times: {times}')\n",
572 | " return image\n",
573 | "\n",
574 | "\n",
575 | "def is_image_file(filename):\n",
576 | " \"\"\"Check if it is an valid image file by extension.\"\"\"\n",
577 | " return any(\n",
578 | " filename.endswith(extension)\n",
579 | " for extension in ['jpeg', 'JPEG', 'jpg', 'png', 'JPG', 'PNG', 'gif'])\n",
580 | "\n",
581 | "\n",
582 | "def save_img(img, pth):\n",
583 | " \"\"\"Save an image to disk.\n",
584 | "\n",
585 | " Args:\n",
586 | " img: jnp.ndarry, [height, width, channels], img will be clipped to [0, 1]\n",
587 | " before saved to pth.\n",
588 | " pth: string, path to save the image to.\n",
589 | " \"\"\"\n",
590 | " Image.fromarray(np.array(\n",
591 | " (np.clip(img, 0., 1.) * 255.).astype(jnp.uint8))).save(pth, 'PNG')\n",
592 | "\n",
593 | "\n",
594 | "def make_shape_even(image):\n",
595 | " \"\"\"Pad the image to have even shapes.\"\"\"\n",
596 | " height, width = image.shape[0], image.shape[1]\n",
597 | " padh = 1 if height % 2 != 0 else 0\n",
598 | " padw = 1 if width % 2 != 0 else 0\n",
599 | " image = jnp.pad(image, [(0, padh), (0, padw), (0, 0)], mode='reflect')\n",
600 | " return image\n",
601 | "\n",
602 | "\n",
603 | "# Refactored code --------------------------------------------------------------------------------------------------------------------\n",
604 | "\n",
605 | "def build_model(task = \"Enhancement\"):\n",
606 | " model_mod = importlib.import_module(f'maxim.models.{_MODEL_FILENAME}')\n",
607 | " model_configs = ml_collections.ConfigDict(_MODEL_CONFIGS)\n",
608 | "\n",
609 | " model_configs.variant = _MODEL_VARIANT_DICT[task]\n",
610 | "\n",
611 | " model = model_mod.Model(**model_configs)\n",
612 | " return model\n",
613 | "\n",
614 | "\n",
615 | "def pre_process(input_file):\n",
616 | " '''\n",
617 | " Pre-process the image before sending to the model\n",
618 | " '''\n",
619 | " input_img = np.asarray(Image.open(input_file).convert('RGB'),np.float32) / 255.\n",
620 | " # Padding images to have even shapes\n",
621 | " height, width = input_img.shape[0], input_img.shape[1]\n",
622 | " input_img = make_shape_even(input_img)\n",
623 | " height_even, width_even = input_img.shape[0], input_img.shape[1]\n",
624 | "\n",
625 | " # padding images to be multiplies of 64\n",
626 | " input_img = mod_padding_symmetric(input_img, factor=64)\n",
627 | " input_img = np.expand_dims(input_img, axis=0)\n",
628 | "\n",
629 | " return input_img, height, width, height_even, width_even\n",
630 | "\n",
631 | "\n",
632 | "def predict(input_img):\n",
633 | " # handle multi-stage outputs, obtain the last scale output of last stage\n",
634 | " return model.apply({'params': flax.core.freeze(params)}, input_img)\n",
635 | "\n",
636 | "\n",
637 | "def post_process(preds, height, width, height_even, width_even):\n",
638 | " '''\n",
639 | " Post process the image coming out from prediction\n",
640 | " '''\n",
641 | " if isinstance(preds, list):\n",
642 | " preds = preds[-1]\n",
643 | " if isinstance(preds, list):\n",
644 | " preds = preds[-1]\n",
645 | "\n",
646 | " # De-ensemble by averaging inferenced results.\n",
647 | " preds = np.array(preds[0], np.float32)\n",
648 | "\n",
649 | " # unpad images to get the original resolution\n",
650 | " new_height, new_width = preds.shape[0], preds.shape[1]\n",
651 | " h_start = new_height // 2 - height_even // 2\n",
652 | " h_end = h_start + height\n",
653 | " w_start = new_width // 2 - width_even // 2\n",
654 | " w_end = w_start + width\n",
655 | " preds = preds[h_start:h_end, w_start:w_end, :]\n",
656 | " return np.array((np.clip(preds, 0., 1.) * 255.).astype(jnp.uint8))"
657 | ]
658 | },
659 | {
660 | "cell_type": "markdown",
661 | "metadata": {
662 | "id": "JOfe8u7_Wxks"
663 | },
664 | "source": [
665 | "# Default Configs and Model Building\n",
666 | "**Steps**:\n",
667 | "1. Get the name of `task` and the respective `ckpt` (pre-trained saved model for that task) [Follow this link for task name and model](https://github.com/google-research/maxim#results-and-pre-trained-models)\n",
668 | "2. Pass in the proper `task` and `ckpt_path` to the `DummyFlags`\n",
669 | "3. Build Model"
670 | ]
671 | },
672 | {
673 | "cell_type": "code",
674 | "execution_count": null,
675 | "metadata": {
676 | "id": "Fcp68HNFf2Fy",
677 | "colab": {
678 | "base_uri": "https://localhost:8080/"
679 | },
680 | "outputId": "09761106-e8cd-4880-bea6-e6e0840effff"
681 | },
682 | "outputs": [
683 | {
684 | "output_type": "stream",
685 | "name": "stderr",
686 | "text": [
687 | "Downloading...\n",
688 | "From: https://drive.google.com/uc?id=1-BRKozXh81PtwoMZ9QN3kCAieLzozHIq\n",
689 | "To: /content/maxim/adobe.npz\n",
690 | "100%|██████████| 172M/172M [00:01<00:00, 166MB/s]\n"
691 | ]
692 | }
693 | ],
694 | "source": [
695 | "weight_drive_path = 'https://drive.google.com/uc?id=1-BRKozXh81PtwoMZ9QN3kCAieLzozHIq' # Path of the weights file which in the Google Drive\n",
696 | "MODEL_PATH = './adobe.npz' # name of the model to be saved as\n",
697 | "\n",
698 | "gdown.download(weight_drive_path, MODEL_PATH, quiet=False) # Download Model weights to your current instance\n",
699 | "\n",
700 | "\n",
701 | "FLAGS = DummyFlags(ckpt_path = MODEL_PATH, task = \"Enhancement\") # Path to your checkpoint and task name\n",
702 | "\n",
703 | "params = get_params(FLAGS.ckpt_path) # Parse the config\n",
704 | "\n",
705 | "model = build_model() # Build Model"
706 | ]
707 | },
708 | {
709 | "cell_type": "markdown",
710 | "metadata": {
711 | "id": "c_zxNUE4TugU"
712 | },
713 | "source": [
714 | "# Inference\n",
715 | "For Inference, you just need to pasd the *Image Path* to the the `predict` function. Result will be a `Numpy` array. You can easily save that by converting to `PIL` image.\n",
716 | "\n",
717 | "\n",
718 | "**Note**: You might get `OOM` or Out of memory issue which is not a big deal if you image size is too big. In that case, you just need to use the `resize` function\n",
719 | "\n"
720 | ]
721 | },
722 | {
723 | "cell_type": "code",
724 | "execution_count": null,
725 | "metadata": {
726 | "id": "jC35dD6ViBwZ"
727 | },
728 | "outputs": [],
729 | "source": [
730 | "# image_path = \"path/to/my/image.extension\" # your image path\n",
731 | "# enhanced_image_array = predict(image_path) # Get predictions\n",
732 | "\n",
733 | "# enhanced_pil_image = Image.fromarray(enhanced_image_array) # get PIL image from array\n",
734 | "# enhanced_pil_image.save(\"path/to/output/directory/image.extension\") # Save the image\n"
735 | ]
736 | },
737 | {
738 | "cell_type": "markdown",
739 | "metadata": {
740 | "id": "7x5dCbuOvIDz"
741 | },
742 | "source": [
743 | "# Test Images from Drive and Save\n",
744 | "\n",
745 | "**Note**: For huge number of images (say 50 or more), copy all the images from Google Drive to the current machine's drive else it will make the process so slow. And also for saving the enhanced image to drive, Get predictions for all the images at once, Save them here first and them copy a zip file to the drive."
746 | ]
747 | },
748 | {
749 | "cell_type": "code",
750 | "execution_count": null,
751 | "metadata": {
752 | "id": "W3YwEUHmtQSl"
753 | },
754 | "outputs": [],
755 | "source": [
756 | "# images = [\"../gdrive/My Drive/maxim/input/\"+i for i in os.listdir(\"../gdrive/My Drive/maxim/input/\") if i.endswith(('jpeg', 'png', 'jpg',\"PNG\",\"JPEG\",\"JPG\"))]\n",
757 | "\n",
758 | "# # _ = [resize(path, 1920, save_image=True) for path in images] # Resize Images to 1920 as the max dimension's size else it'll blow the GPU / CPU memory\n",
759 | "\n",
760 | "\n",
761 | "# for path in images:\n",
762 | "# im = Image.fromarray(predict(path))\n",
763 | "# im.save(\"../gdrive/My Drive/maxim/output/\"+path.split('/')[-1])\n",
764 | "\n"
765 | ]
766 | },
767 | {
768 | "cell_type": "markdown",
769 | "metadata": {
770 | "id": "f9Wv9yIIamJL"
771 | },
772 | "source": [
773 | "# Visualization\n",
774 | "\n",
775 | "The below code demonstrates how to predict from Image URL. You can directly use `predict(image_path)`"
776 | ]
777 | },
778 | {
779 | "cell_type": "code",
780 | "execution_count": null,
781 | "metadata": {
782 | "id": "Q1T544sZan4d",
783 | "colab": {
784 | "base_uri": "https://localhost:8080/",
785 | "height": 491
786 | },
787 | "outputId": "17f4b919-c811-4659-a178-95d02d8660dc"
788 | },
789 | "outputs": [
790 | {
791 | "output_type": "stream",
792 | "text": [
793 | "\u001b[1;30;43mThis cell output is too large and can only be displayed while logged in.\u001b[0m\n"
794 | ]
795 | }
796 | ],
797 | "source": [
798 | "import requests\n",
799 | "from io import BytesIO\n",
800 | "\n",
801 | "url = \"https://phototraces.b-cdn.net/wp-content/uploads/2021/02/id_Free_RAW_Photos_for_Editing_09_Uneditedd.jpg\"\n",
802 | "# url = \"https://phototraces.b-cdn.net/wp-content/uploads/2021/03/Free_RAW_Photos_for_Editing_13_Unedited.jpg\"\n",
803 | "\n",
804 | "image_bytes = BytesIO(requests.get(url).content)\n",
805 | "\n",
806 | "input_img, height, width, height_even, width_even = pre_process(image_bytes)\n",
807 | "preds = predict(input_img)\n",
808 | "result = post_process(preds, height, width, height_even, width_even)\n",
809 | "\n",
810 | "f, ax = plt.subplots(1,2, figsize = (35,20))\n",
811 | "\n",
812 | "ax[0].imshow(np.array(Image.open(image_bytes))) # Original image\n",
813 | "ax[1].imshow(result) # retouched image\n",
814 | "\n",
815 | "ax[0].set_title(\"Original Image\")\n",
816 | "ax[1].set_title(\"Enhanced Image\")\n",
817 | "\n",
818 | "plt.show()"
819 | ]
820 | },
821 | {
822 | "cell_type": "code",
823 | "source": [
824 | "tf_predict = tf.function(\n",
825 | " jax2tf.convert(predict, enable_xla=False),\n",
826 | " input_signature=[\n",
827 | " tf.TensorSpec(shape=[1, 704, 1024, 3], dtype=tf.float32, name='input_image')\n",
828 | " ],\n",
829 | " autograph=False)"
830 | ],
831 | "metadata": {
832 | "id": "nZGyNC8pttv7"
833 | },
834 | "execution_count": null,
835 | "outputs": []
836 | },
837 | {
838 | "cell_type": "code",
839 | "source": [
840 | "converter = tf.lite.TFLiteConverter.from_concrete_functions(\n",
841 | " [tf_predict.get_concrete_function()], tf_predict)\n",
842 | "\n",
843 | "converter.target_spec.supported_ops = [\n",
844 | " tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.\n",
845 | " tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.\n",
846 | "]\n",
847 | "tflite_float_model = converter.convert()\n",
848 | "\n",
849 | "with open('./float_model.tflite', \"wb\") as f: f.write(tflite_float_model)"
850 | ],
851 | "metadata": {
852 | "colab": {
853 | "base_uri": "https://localhost:8080/"
854 | },
855 | "id": "igxeGKfIyqLS",
856 | "outputId": "2c7d6ae3-8c84-421f-dc83-f6e0c69cc6c5"
857 | },
858 | "execution_count": null,
859 | "outputs": [
860 | {
861 | "output_type": "stream",
862 | "name": "stdout",
863 | "text": [
864 | "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n"
865 | ]
866 | },
867 | {
868 | "output_type": "stream",
869 | "name": "stderr",
870 | "text": [
871 | "WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.\n",
872 | "WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded\n"
873 | ]
874 | }
875 | ]
876 | },
877 | {
878 | "cell_type": "code",
879 | "source": [
880 | "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
881 | "tflite_quantized_model = converter.convert()\n",
882 | "\n",
883 | "with open('./quantized.tflite', 'wb') as f: f.write(tflite_quantized_model)"
884 | ],
885 | "metadata": {
886 | "colab": {
887 | "base_uri": "https://localhost:8080/"
888 | },
889 | "id": "pTbjUpn-yxl1",
890 | "outputId": "fcb29b97-825c-4d6b-b921-0452bcf003a0"
891 | },
892 | "execution_count": null,
893 | "outputs": [
894 | {
895 | "output_type": "stream",
896 | "name": "stderr",
897 | "text": [
898 | "WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded\n"
899 | ]
900 | }
901 | ]
902 | },
903 | {
904 | "cell_type": "code",
905 | "source": [
906 | "# Load quantized TFLite model\n",
907 | "tflite_interpreter_quant = tf.lite.Interpreter(model_path='./maxim/quantized.tflite')\n",
908 | "\n",
909 | "# Learn about its input and output details\n",
910 | "input_details = tflite_interpreter_quant.get_input_details()\n",
911 | "output_details = tflite_interpreter_quant.get_output_details()\n",
912 | "\n",
913 | "# Resize input and output tensors to handle batch of desired size\n",
914 | "# tflite_interpreter_quant.resize_tensor_input(input_details[0]['index'], (1, 704, 1024, 3))\n",
915 | "# tflite_interpreter_quant.resize_tensor_input(output_details[0]['index'], (1, 176, 256, 3))\n",
916 | "tflite_interpreter_quant.allocate_tensors()\n",
917 | "\n",
918 | "input_details = tflite_interpreter_quant.get_input_details()\n",
919 | "output_details = tflite_interpreter_quant.get_output_details()\n",
920 | "\n",
921 | "\n",
922 | "# # Run inference\n",
923 | "val_image_batch = tf.random.normal(shape = (1, 704, 1024, 3), dtype = tf.float32)\n",
924 | "tflite_interpreter_quant.set_tensor(input_details[0]['index'], val_image_batch)\n",
925 | "\n",
926 | "tflite_interpreter_quant.invoke()\n",
927 | "\n",
928 | "tflite_q_model_predictions = tflite_interpreter_quant.get_tensor(output_details[0]['index'])\n",
929 | "print(\"\\nPrediction results shape:\", tflite_q_model_predictions.shape)"
930 | ],
931 | "metadata": {
932 | "id": "XY0M5SFg4Zw8"
933 | },
934 | "execution_count": null,
935 | "outputs": []
936 | },
937 | {
938 | "cell_type": "code",
939 | "source": [],
940 | "metadata": {
941 | "id": "RgsmrC0U7v5h"
942 | },
943 | "execution_count": null,
944 | "outputs": []
945 | }
946 | ],
947 | "metadata": {
948 | "colab": {
949 | "collapsed_sections": [],
950 | "provenance": []
951 | },
952 | "kernelspec": {
953 | "display_name": "Python 3",
954 | "name": "python3"
955 | },
956 | "language_info": {
957 | "name": "python"
958 | },
959 | "gpuClass": "standard"
960 | },
961 | "nbformat": 4,
962 | "nbformat_minor": 0
963 | }
--------------------------------------------------------------------------------
/maxim/__init__.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
--------------------------------------------------------------------------------
/maxim/images/Deblurring/input/109fromGOPR1096.MP4.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deblurring/input/109fromGOPR1096.MP4.png
--------------------------------------------------------------------------------
/maxim/images/Deblurring/input/110fromGOPR1087.MP4.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deblurring/input/110fromGOPR1087.MP4.png
--------------------------------------------------------------------------------
/maxim/images/Deblurring/input/1fromGOPR0950.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deblurring/input/1fromGOPR0950.png
--------------------------------------------------------------------------------
/maxim/images/Deblurring/input/1fromGOPR1096.MP4.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deblurring/input/1fromGOPR1096.MP4.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/0003_0.8_0.2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/0003_0.8_0.2.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/0010_0.95_0.16.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/0010_0.95_0.16.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/0014_0.8_0.12.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/0014_0.8_0.12.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/0048_0.9_0.2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/0048_0.9_0.2.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/1440_10.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/1440_10.png
--------------------------------------------------------------------------------
/maxim/images/Dehazing/input/1444_10.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Dehazing/input/1444_10.png
--------------------------------------------------------------------------------
/maxim/images/Denoising/input/0003_30.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Denoising/input/0003_30.png
--------------------------------------------------------------------------------
/maxim/images/Denoising/input/0011_23.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Denoising/input/0011_23.png
--------------------------------------------------------------------------------
/maxim/images/Denoising/input/0013_19.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Denoising/input/0013_19.png
--------------------------------------------------------------------------------
/maxim/images/Denoising/input/0039_04.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Denoising/input/0039_04.png
--------------------------------------------------------------------------------
/maxim/images/Deraining/input/0.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deraining/input/0.jpg
--------------------------------------------------------------------------------
/maxim/images/Deraining/input/1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deraining/input/1.png
--------------------------------------------------------------------------------
/maxim/images/Deraining/input/15.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deraining/input/15.png
--------------------------------------------------------------------------------
/maxim/images/Deraining/input/55.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Deraining/input/55.png
--------------------------------------------------------------------------------
/maxim/images/Enhancement/input/1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Enhancement/input/1.png
--------------------------------------------------------------------------------
/maxim/images/Enhancement/input/111.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Enhancement/input/111.png
--------------------------------------------------------------------------------
/maxim/images/Enhancement/input/748.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Enhancement/input/748.png
--------------------------------------------------------------------------------
/maxim/images/Enhancement/input/a4541-DSC_0040-2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/Enhancement/input/a4541-DSC_0040-2.png
--------------------------------------------------------------------------------
/maxim/images/overview.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/google-research/maxim/3c8265171ffccc80c3c9124844aef0d381609956/maxim/images/overview.png
--------------------------------------------------------------------------------
/maxim/models/__init__.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
--------------------------------------------------------------------------------
/maxim/models/maxim.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | """Main file for the MAXIM model."""
16 |
17 | import functools
18 | from typing import Any, Sequence, Tuple
19 |
20 | import einops
21 | import flax.linen as nn
22 | import jax
23 | import jax.numpy as jnp
24 |
25 |
26 | Conv3x3 = functools.partial(nn.Conv, kernel_size=(3, 3))
27 | Conv1x1 = functools.partial(nn.Conv, kernel_size=(1, 1))
28 | ConvT_up = functools.partial(nn.ConvTranspose,
29 | kernel_size=(2, 2),
30 | strides=(2, 2))
31 | Conv_down = functools.partial(nn.Conv,
32 | kernel_size=(4, 4),
33 | strides=(2, 2))
34 |
35 | weight_initializer = nn.initializers.normal(stddev=2e-2)
36 |
37 |
38 | class MlpBlock(nn.Module):
39 | """A 1-hidden-layer MLP block, applied over the last dimension."""
40 | mlp_dim: int
41 | dropout_rate: float = 0.0
42 | use_bias: bool = True
43 |
44 | @nn.compact
45 | def __call__(self, x, deterministic=True):
46 | n, h, w, d = x.shape
47 | x = nn.Dense(self.mlp_dim, use_bias=self.use_bias,
48 | kernel_init=weight_initializer)(x)
49 | x = nn.gelu(x)
50 | x = nn.Dropout(rate=self.dropout_rate)(x, deterministic)
51 | x = nn.Dense(d, use_bias=self.use_bias,
52 | kernel_init=weight_initializer)(x)
53 | return x
54 |
55 |
56 | def block_images_einops(x, patch_size):
57 | """Image to patches."""
58 | batch, height, width, channels = x.shape
59 | grid_height = height // patch_size[0]
60 | grid_width = width // patch_size[1]
61 | x = einops.rearrange(
62 | x, "n (gh fh) (gw fw) c -> n (gh gw) (fh fw) c",
63 | gh=grid_height, gw=grid_width, fh=patch_size[0], fw=patch_size[1])
64 | return x
65 |
66 |
67 | def unblock_images_einops(x, grid_size, patch_size):
68 | """patches to images."""
69 | x = einops.rearrange(
70 | x, "n (gh gw) (fh fw) c -> n (gh fh) (gw fw) c",
71 | gh=grid_size[0], gw=grid_size[1], fh=patch_size[0], fw=patch_size[1])
72 | return x
73 |
74 |
75 | class UpSampleRatio(nn.Module):
76 | """Upsample features given a ratio > 0."""
77 | features: int
78 | ratio: float
79 | use_bias: bool = True
80 |
81 | @nn.compact
82 | def __call__(self, x):
83 | n, h, w, c = x.shape
84 | x = jax.image.resize(
85 | x,
86 | shape=(n, int(h * self.ratio), int(w * self.ratio), c),
87 | method="bilinear")
88 | x = Conv1x1(features=self.features, use_bias=self.use_bias)(x)
89 | return x
90 |
91 |
92 | class CALayer(nn.Module):
93 | """Squeeze-and-excitation block for channel attention.
94 |
95 | ref: https://arxiv.org/abs/1709.01507
96 | """
97 | features: int
98 | reduction: int = 4
99 | use_bias: bool = True
100 |
101 | @nn.compact
102 | def __call__(self, x):
103 | # 2D global average pooling
104 | y = jnp.mean(x, axis=[1, 2], keepdims=True)
105 | # Squeeze (in Squeeze-Excitation)
106 | y = Conv1x1(self.features // self.reduction, use_bias=self.use_bias)(y)
107 | y = nn.relu(y)
108 | # Excitation (in Squeeze-Excitation)
109 | y = Conv1x1(self.features, use_bias=self.use_bias)(y)
110 | y = nn.sigmoid(y)
111 | return x * y
112 |
113 |
114 | class RCAB(nn.Module):
115 | """Residual channel attention block. Contains LN,Conv,lRelu,Conv,SELayer."""
116 | features: int
117 | reduction: int = 4
118 | lrelu_slope: float = 0.2
119 | use_bias: bool = True
120 |
121 | @nn.compact
122 | def __call__(self, x):
123 | shortcut = x
124 | x = nn.LayerNorm(name="LayerNorm")(x)
125 | x = Conv3x3(features=self.features, use_bias=self.use_bias, name="conv1")(x)
126 | x = nn.leaky_relu(x, negative_slope=self.lrelu_slope)
127 | x = Conv3x3(features=self.features, use_bias=self.use_bias, name="conv2")(x)
128 | x = CALayer(features=self.features, reduction=self.reduction,
129 | use_bias=self.use_bias, name="channel_attention")(x)
130 | return x + shortcut
131 |
132 |
133 | class GridGatingUnit(nn.Module):
134 | """A SpatialGatingUnit as defined in the gMLP paper.
135 |
136 | The 'spatial' dim is defined as the second last.
137 | If applied on other dims, you should swapaxes first.
138 | """
139 | use_bias: bool = True
140 |
141 | @nn.compact
142 | def __call__(self, x):
143 | u, v = jnp.split(x, 2, axis=-1)
144 | v = nn.LayerNorm(name="intermediate_layernorm")(v)
145 | n = x.shape[-3] # get spatial dim
146 | v = jnp.swapaxes(v, -1, -3)
147 | v = nn.Dense(n, use_bias=self.use_bias, kernel_init=weight_initializer)(v)
148 | v = jnp.swapaxes(v, -1, -3)
149 | return u * (v + 1.)
150 |
151 |
152 | class GridGmlpLayer(nn.Module):
153 | """Grid gMLP layer that performs global mixing of tokens."""
154 | grid_size: Sequence[int]
155 | use_bias: bool = True
156 | factor: int = 2
157 | dropout_rate: float = 0.0
158 |
159 | @nn.compact
160 | def __call__(self, x, deterministic=True):
161 | n, h, w, num_channels = x.shape
162 | gh, gw = self.grid_size
163 | fh, fw = h // gh, w // gw
164 | x = block_images_einops(x, patch_size=(fh, fw))
165 | # gMLP1: Global (grid) mixing part, provides global grid communication.
166 | y = nn.LayerNorm(name="LayerNorm")(x)
167 | y = nn.Dense(num_channels * self.factor, use_bias=self.use_bias,
168 | kernel_init=weight_initializer, name="in_project")(y)
169 | y = nn.gelu(y)
170 | y = GridGatingUnit(use_bias=self.use_bias, name="GridGatingUnit")(y)
171 | y = nn.Dense(num_channels, use_bias=self.use_bias,
172 | kernel_init=weight_initializer, name="out_project")(y)
173 | y = nn.Dropout(self.dropout_rate)(y, deterministic)
174 | x = x + y
175 | x = unblock_images_einops(x, grid_size=(gh, gw), patch_size=(fh, fw))
176 | return x
177 |
178 |
179 | class BlockGatingUnit(nn.Module):
180 | """A SpatialGatingUnit as defined in the gMLP paper.
181 |
182 | The 'spatial' dim is defined as the **second last**.
183 | If applied on other dims, you should swapaxes first.
184 | """
185 | use_bias: bool = True
186 |
187 | @nn.compact
188 | def __call__(self, x):
189 | u, v = jnp.split(x, 2, axis=-1)
190 | v = nn.LayerNorm(name="intermediate_layernorm")(v)
191 | n = x.shape[-2] # get spatial dim
192 | v = jnp.swapaxes(v, -1, -2)
193 | v = nn.Dense(n, use_bias=self.use_bias, kernel_init=weight_initializer)(v)
194 | v = jnp.swapaxes(v, -1, -2)
195 | return u * (v + 1.)
196 |
197 |
198 | class BlockGmlpLayer(nn.Module):
199 | """Block gMLP layer that performs local mixing of tokens."""
200 | block_size: Sequence[int]
201 | use_bias: bool = True
202 | factor: int = 2
203 | dropout_rate: float = 0.0
204 |
205 | @nn.compact
206 | def __call__(self, x, deterministic=True):
207 | n, h, w, num_channels = x.shape
208 | fh, fw = self.block_size
209 | gh, gw = h // fh, w // fw
210 | x = block_images_einops(x, patch_size=(fh, fw))
211 | # MLP2: Local (block) mixing part, provides within-block communication.
212 | y = nn.LayerNorm(name="LayerNorm")(x)
213 | y = nn.Dense(num_channels * self.factor, use_bias=self.use_bias,
214 | kernel_init=weight_initializer, name="in_project")(y)
215 | y = nn.gelu(y)
216 | y = BlockGatingUnit(use_bias=self.use_bias, name="BlockGatingUnit")(y)
217 | y = nn.Dense(num_channels, use_bias=self.use_bias,
218 | kernel_init=weight_initializer, name="out_project")(y)
219 | y = nn.Dropout(self.dropout_rate)(y, deterministic)
220 | x = x + y
221 | x = unblock_images_einops(x, grid_size=(gh, gw), patch_size=(fh, fw))
222 | return x
223 |
224 |
225 | class ResidualSplitHeadMultiAxisGmlpLayer(nn.Module):
226 | """The multi-axis gated MLP block."""
227 | block_size: Sequence[int]
228 | grid_size: Sequence[int]
229 | block_gmlp_factor: int = 2
230 | grid_gmlp_factor: int = 2
231 | input_proj_factor: int = 2
232 | use_bias: bool = True
233 | dropout_rate: float = 0.0
234 |
235 | @nn.compact
236 | def __call__(self, x, deterministic=True):
237 | shortcut = x
238 | n, h, w, num_channels = x.shape
239 | x = nn.LayerNorm(name="LayerNorm_in")(x)
240 | x = nn.Dense(num_channels * self.input_proj_factor, use_bias=self.use_bias,
241 | kernel_init=weight_initializer, name="in_project")(x)
242 | x = nn.gelu(x)
243 |
244 | u, v = jnp.split(x, 2, axis=-1)
245 | # GridGMLPLayer
246 | u = GridGmlpLayer(
247 | grid_size=self.grid_size,
248 | factor=self.grid_gmlp_factor,
249 | use_bias=self.use_bias,
250 | dropout_rate=self.dropout_rate,
251 | name="GridGmlpLayer")(u, deterministic)
252 |
253 | # BlockGMLPLayer
254 | v = BlockGmlpLayer(
255 | block_size=self.block_size,
256 | factor=self.block_gmlp_factor,
257 | use_bias=self.use_bias,
258 | dropout_rate=self.dropout_rate,
259 | name="BlockGmlpLayer")(v, deterministic)
260 |
261 | x = jnp.concatenate([u, v], axis=-1)
262 |
263 | x = nn.Dense(num_channels, use_bias=self.use_bias,
264 | kernel_init=weight_initializer, name="out_project")(x)
265 | x = nn.Dropout(self.dropout_rate)(x, deterministic)
266 | x = x + shortcut
267 | return x
268 |
269 |
270 | class RDCAB(nn.Module):
271 | """Residual dense channel attention block. Used in Bottlenecks."""
272 | features: int
273 | reduction: int = 16
274 | use_bias: bool = True
275 | dropout_rate: float = 0.0
276 |
277 | @nn.compact
278 | def __call__(self, x, deterministic=True):
279 | y = nn.LayerNorm(name="LayerNorm")(x)
280 | y = MlpBlock(
281 | mlp_dim=self.features,
282 | dropout_rate=self.dropout_rate,
283 | use_bias=self.use_bias,
284 | name="channel_mixing")(
285 | y, deterministic=deterministic)
286 | y = CALayer(
287 | features=self.features,
288 | reduction=self.reduction,
289 | use_bias=self.use_bias,
290 | name="channel_attention")(
291 | y)
292 | x = x + y
293 | return x
294 |
295 |
296 | class BottleneckBlock(nn.Module):
297 | """The bottleneck block consisting of multi-axis gMLP block and RDCAB."""
298 | features: int
299 | block_size: Sequence[int]
300 | grid_size: Sequence[int]
301 | num_groups: int = 1
302 | block_gmlp_factor: int = 2
303 | grid_gmlp_factor: int = 2
304 | input_proj_factor: int = 2
305 | channels_reduction: int = 4
306 | dropout_rate: float = 0.0
307 | use_bias: bool = True
308 |
309 | @nn.compact
310 | def __call__(self, x, deterministic):
311 | """Applies the Mixer block to inputs."""
312 | assert x.ndim == 4 # Input has shape [batch, h, w, c]
313 | n, h, w, num_channels = x.shape
314 |
315 | # input projection
316 | x = Conv1x1(self.features, use_bias=self.use_bias, name="input_proj")(x)
317 | shortcut_long = x
318 |
319 | for i in range(self.num_groups):
320 | x = ResidualSplitHeadMultiAxisGmlpLayer(
321 | grid_size=self.grid_size,
322 | block_size=self.block_size,
323 | grid_gmlp_factor=self.grid_gmlp_factor,
324 | block_gmlp_factor=self.block_gmlp_factor,
325 | input_proj_factor=self.input_proj_factor,
326 | use_bias=self.use_bias,
327 | dropout_rate=self.dropout_rate,
328 | name=f"SplitHeadMultiAxisGmlpLayer_{i}")(x, deterministic)
329 | # Channel-mixing part, which provides within-patch communication.
330 | x = RDCAB(
331 | features=self.features,
332 | reduction=self.channels_reduction,
333 | use_bias=self.use_bias,
334 | name=f"channel_attention_block_1_{i}")(
335 | x)
336 |
337 | # long skip-connect
338 | x = x + shortcut_long
339 | return x
340 |
341 |
342 | class UNetEncoderBlock(nn.Module):
343 | """Encoder block in MAXIM."""
344 | features: int
345 | block_size: Sequence[int]
346 | grid_size: Sequence[int]
347 | num_groups: int = 1
348 | lrelu_slope: float = 0.2
349 | block_gmlp_factor: int = 2
350 | grid_gmlp_factor: int = 2
351 | input_proj_factor: int = 2
352 | channels_reduction: int = 4
353 | dropout_rate: float = 0.0
354 | downsample: bool = True
355 | use_global_mlp: bool = True
356 | use_bias: bool = True
357 | use_cross_gating: bool = False
358 |
359 | @nn.compact
360 | def __call__(self, x: jnp.ndarray, skip: jnp.ndarray = None,
361 | enc: jnp.ndarray = None, dec: jnp.ndarray = None, *,
362 | deterministic: bool = True) -> jnp.ndarray:
363 | if skip is not None:
364 | x = jnp.concatenate([x, skip], axis=-1)
365 |
366 | # convolution-in
367 | x = Conv1x1(self.features, use_bias=self.use_bias)(x)
368 | shortcut_long = x
369 |
370 | for i in range(self.num_groups):
371 | if self.use_global_mlp:
372 | x = ResidualSplitHeadMultiAxisGmlpLayer(
373 | grid_size=self.grid_size,
374 | block_size=self.block_size,
375 | grid_gmlp_factor=self.grid_gmlp_factor,
376 | block_gmlp_factor=self.block_gmlp_factor,
377 | input_proj_factor=self.input_proj_factor,
378 | use_bias=self.use_bias,
379 | dropout_rate=self.dropout_rate,
380 | name=f"SplitHeadMultiAxisGmlpLayer_{i}")(x, deterministic)
381 | x = RCAB(
382 | features=self.features,
383 | reduction=self.channels_reduction,
384 | use_bias=self.use_bias,
385 | name=f"channel_attention_block_1{i}")(x)
386 |
387 | x = x + shortcut_long
388 |
389 | if enc is not None and dec is not None:
390 | assert self.use_cross_gating
391 | x, _ = CrossGatingBlock(
392 | features=self.features,
393 | block_size=self.block_size,
394 | grid_size=self.grid_size,
395 | dropout_rate=self.dropout_rate,
396 | input_proj_factor=self.input_proj_factor,
397 | upsample_y=False,
398 | use_bias=self.use_bias,
399 | name="cross_gating_block")(
400 | x, enc + dec, deterministic=deterministic)
401 |
402 | if self.downsample:
403 | x_down = Conv_down(self.features, use_bias=self.use_bias)(x)
404 | return x_down, x
405 | else:
406 | return x
407 |
408 |
409 | class UNetDecoderBlock(nn.Module):
410 | """Decoder block in MAXIM."""
411 | features: int
412 | block_size: Sequence[int]
413 | grid_size: Sequence[int]
414 | num_groups: int = 1
415 | lrelu_slope: float = 0.2
416 | block_gmlp_factor: int = 2
417 | grid_gmlp_factor: int = 2
418 | input_proj_factor: int = 2
419 | channels_reduction: int = 4
420 | dropout_rate: float = 0.0
421 | downsample: bool = True
422 | use_global_mlp: bool = True
423 | use_bias: bool = True
424 |
425 | @nn.compact
426 | def __call__(self, x: jnp.ndarray, bridge: jnp.ndarray = None,
427 | deterministic: bool = True) -> jnp.ndarray:
428 | x = ConvT_up(self.features, use_bias=self.use_bias)(x)
429 |
430 | x = UNetEncoderBlock(
431 | self.features,
432 | num_groups=self.num_groups,
433 | lrelu_slope=self.lrelu_slope,
434 | block_size=self.block_size,
435 | grid_size=self.grid_size,
436 | block_gmlp_factor=self.block_gmlp_factor,
437 | grid_gmlp_factor=self.grid_gmlp_factor,
438 | channels_reduction=self.channels_reduction,
439 | use_global_mlp=self.use_global_mlp,
440 | dropout_rate=self.dropout_rate,
441 | downsample=False,
442 | use_bias=self.use_bias)(x, skip=bridge, deterministic=deterministic)
443 | return x
444 |
445 |
446 | class GetSpatialGatingWeights(nn.Module):
447 | """Get gating weights for cross-gating MLP block."""
448 | features: int
449 | block_size: Sequence[int]
450 | grid_size: Sequence[int]
451 | input_proj_factor: int = 2
452 | dropout_rate: float = 0.0
453 | use_bias: bool = True
454 |
455 | @nn.compact
456 | def __call__(self, x, deterministic):
457 | n, h, w, num_channels = x.shape
458 |
459 | # input projection
460 | x = nn.LayerNorm(name="LayerNorm_in")(x)
461 | x = nn.Dense(
462 | num_channels * self.input_proj_factor,
463 | use_bias=self.use_bias,
464 | name="in_project")(
465 | x)
466 | x = nn.gelu(x)
467 | u, v = jnp.split(x, 2, axis=-1)
468 |
469 | # Get grid MLP weights
470 | gh, gw = self.grid_size
471 | fh, fw = h // gh, w // gw
472 | u = block_images_einops(u, patch_size=(fh, fw))
473 | dim_u = u.shape[-3]
474 | u = jnp.swapaxes(u, -1, -3)
475 | u = nn.Dense(
476 | dim_u, use_bias=self.use_bias, kernel_init=nn.initializers.normal(2e-2),
477 | bias_init=nn.initializers.ones)(u)
478 | u = jnp.swapaxes(u, -1, -3)
479 | u = unblock_images_einops(u, grid_size=(gh, gw), patch_size=(fh, fw))
480 |
481 | # Get Block MLP weights
482 | fh, fw = self.block_size
483 | gh, gw = h // fh, w // fw
484 | v = block_images_einops(v, patch_size=(fh, fw))
485 | dim_v = v.shape[-2]
486 | v = jnp.swapaxes(v, -1, -2)
487 | v = nn.Dense(
488 | dim_v, use_bias=self.use_bias, kernel_init=nn.initializers.normal(2e-2),
489 | bias_init=nn.initializers.ones)(v)
490 | v = jnp.swapaxes(v, -1, -2)
491 | v = unblock_images_einops(v, grid_size=(gh, gw), patch_size=(fh, fw))
492 |
493 | x = jnp.concatenate([u, v], axis=-1)
494 | x = nn.Dense(num_channels, use_bias=self.use_bias, name="out_project")(x)
495 | x = nn.Dropout(self.dropout_rate)(x, deterministic)
496 | return x
497 |
498 |
499 | class CrossGatingBlock(nn.Module):
500 | """Cross-gating MLP block."""
501 | features: int
502 | block_size: Sequence[int]
503 | grid_size: Sequence[int]
504 | dropout_rate: float = 0.0
505 | input_proj_factor: int = 2
506 | upsample_y: bool = True
507 | use_bias: bool = True
508 |
509 | @nn.compact
510 | def __call__(self, x, y, deterministic=True):
511 | # Upscale Y signal, y is the gating signal.
512 | if self.upsample_y:
513 | y = ConvT_up(self.features, use_bias=self.use_bias)(y)
514 |
515 | x = Conv1x1(self.features, use_bias=self.use_bias)(x)
516 | n, h, w, num_channels = x.shape
517 | y = Conv1x1(num_channels, use_bias=self.use_bias)(y)
518 |
519 | assert y.shape == x.shape
520 | shortcut_x = x
521 | shortcut_y = y
522 |
523 | # Get gating weights from X
524 | x = nn.LayerNorm(name="LayerNorm_x")(x)
525 | x = nn.Dense(num_channels, use_bias=self.use_bias, name="in_project_x")(x)
526 | x = nn.gelu(x)
527 | gx = GetSpatialGatingWeights(
528 | features=num_channels,
529 | block_size=self.block_size,
530 | grid_size=self.grid_size,
531 | dropout_rate=self.dropout_rate,
532 | use_bias=self.use_bias,
533 | name="SplitHeadMultiAxisGating_x")(
534 | x, deterministic=deterministic)
535 |
536 | # Get gating weights from Y
537 | y = nn.LayerNorm(name="LayerNorm_y")(y)
538 | y = nn.Dense(num_channels, use_bias=self.use_bias, name="in_project_y")(y)
539 | y = nn.gelu(y)
540 | gy = GetSpatialGatingWeights(
541 | features=num_channels,
542 | block_size=self.block_size,
543 | grid_size=self.grid_size,
544 | dropout_rate=self.dropout_rate,
545 | use_bias=self.use_bias,
546 | name="SplitHeadMultiAxisGating_y")(
547 | y, deterministic=deterministic)
548 |
549 | # Apply cross gating: X = X * GY, Y = Y * GX
550 | y = y * gx
551 | y = nn.Dense(num_channels, use_bias=self.use_bias, name="out_project_y")(y)
552 | y = nn.Dropout(self.dropout_rate)(y, deterministic=deterministic)
553 | y = y + shortcut_y
554 |
555 | x = x * gy # gating x using y
556 | x = nn.Dense(num_channels, use_bias=self.use_bias, name="out_project_x")(x)
557 | x = nn.Dropout(self.dropout_rate)(x, deterministic=deterministic)
558 | x = x + y + shortcut_x # get all aggregated signals
559 | return x, y
560 |
561 |
562 | class SAM(nn.Module):
563 | """Supervised attention module for multi-stage training.
564 |
565 | Introduced by MPRNet [CVPR2021]: https://github.com/swz30/MPRNet
566 | """
567 | features: int
568 | output_channels: int = 3
569 | use_bias: bool = True
570 |
571 | @nn.compact
572 | def __call__(self, x: jnp.ndarray, x_image: jnp.ndarray, *,
573 | train: bool) -> Tuple[jnp.ndarray, jnp.ndarray]:
574 | """Apply the SAM module to the input and features.
575 |
576 | Args:
577 | x: the output features from UNet decoder with shape (h, w, c)
578 | x_image: the input image with shape (h, w, 3)
579 | train: Whether it is training
580 |
581 | Returns:
582 | A tuple of tensors (x1, image) where (x1) is the sam features used for the
583 | next stage, and (image) is the output restored image at current stage.
584 | """
585 | # Get features
586 | x1 = Conv3x3(self.features, use_bias=self.use_bias)(x)
587 |
588 | # Output restored image X_s
589 | if self.output_channels == 3:
590 | image = Conv3x3(self.output_channels, use_bias=self.use_bias)(x) + x_image
591 | else:
592 | image = Conv3x3(self.output_channels, use_bias=self.use_bias)(x)
593 |
594 | # Get attention maps for features
595 | x2 = nn.sigmoid(Conv3x3(self.features, use_bias=self.use_bias)(image))
596 |
597 | # Get attended feature maps
598 | x1 = x1 * x2
599 |
600 | # Residual connection
601 | x1 = x1 + x
602 | return x1, image
603 |
604 |
605 | class MAXIM(nn.Module):
606 | """The MAXIM model function with multi-stage and multi-scale supervision.
607 |
608 | For more model details, please check the CVPR paper:
609 | MAXIM: MUlti-Axis MLP for Image Processing (https://arxiv.org/abs/2201.02973)
610 |
611 | Attributes:
612 | features: initial hidden dimension for the input resolution.
613 | depth: the number of downsampling depth for the model.
614 | num_stages: how many stages to use. It will also affects the output list.
615 | num_groups: how many blocks each stage contains.
616 | use_bias: whether to use bias in all the conv/mlp layers.
617 | num_supervision_scales: the number of desired supervision scales.
618 | lrelu_slope: the negative slope parameter in leaky_relu layers.
619 | use_global_mlp: whether to use the multi-axis gated MLP block (MAB) in each
620 | layer.
621 | use_cross_gating: whether to use the cross-gating MLP block (CGB) in the
622 | skip connections and multi-stage feature fusion layers.
623 | high_res_stages: how many stages are specificied as high-res stages. The
624 | rest (depth - high_res_stages) are called low_res_stages.
625 | block_size_hr: the block_size parameter for high-res stages.
626 | block_size_lr: the block_size parameter for low-res stages.
627 | grid_size_hr: the grid_size parameter for high-res stages.
628 | grid_size_lr: the grid_size parameter for low-res stages.
629 | num_bottleneck_blocks: how many bottleneck blocks.
630 | block_gmlp_factor: the input projection factor for block_gMLP layers.
631 | grid_gmlp_factor: the input projection factor for grid_gMLP layers.
632 | input_proj_factor: the input projection factor for the MAB block.
633 | channels_reduction: the channel reduction factor for SE layer.
634 | num_outputs: the output channels.
635 | dropout_rate: Dropout rate.
636 |
637 | Returns:
638 | The output contains a list of arrays consisting of multi-stage multi-scale
639 | outputs. For example, if num_stages = num_supervision_scales = 3 (the
640 | model used in the paper), the output specs are: outputs =
641 | [[output_stage1_scale1, output_stage1_scale2, output_stage1_scale3],
642 | [output_stage2_scale1, output_stage2_scale2, output_stage2_scale3],
643 | [output_stage3_scale1, output_stage3_scale2, output_stage3_scale3],]
644 | The final output can be retrieved by outputs[-1][-1].
645 | """
646 | features: int = 64
647 | depth: int = 3
648 | num_stages: int = 2
649 | num_groups: int = 1
650 | use_bias: bool = True
651 | num_supervision_scales: int = 1
652 | lrelu_slope: float = 0.2
653 | use_global_mlp: bool = True
654 | use_cross_gating: bool = True
655 | high_res_stages: int = 2
656 | block_size_hr: Sequence[int] = (16, 16)
657 | block_size_lr: Sequence[int] = (8, 8)
658 | grid_size_hr: Sequence[int] = (16, 16)
659 | grid_size_lr: Sequence[int] = (8, 8)
660 | num_bottleneck_blocks: int = 1
661 | block_gmlp_factor: int = 2
662 | grid_gmlp_factor: int = 2
663 | input_proj_factor: int = 2
664 | channels_reduction: int = 4
665 | num_outputs: int = 3
666 | dropout_rate: float = 0.0
667 |
668 | @nn.compact
669 | def __call__(self, x: jnp.ndarray, *, train: bool = False) -> Any:
670 |
671 | n, h, w, c = x.shape # input image shape
672 | shortcuts = []
673 | shortcuts.append(x)
674 | # Get multi-scale input images
675 | for i in range(1, self.num_supervision_scales):
676 | shortcuts.append(jax.image.resize(
677 | x, shape=(n, h // (2**i), w // (2**i), c), method="nearest"))
678 |
679 | # store outputs from all stages and all scales
680 | # Eg, [[(64, 64, 3), (128, 128, 3), (256, 256, 3)], # Stage-1 outputs
681 | # [(64, 64, 3), (128, 128, 3), (256, 256, 3)],] # Stage-2 outputs
682 | outputs_all = []
683 | sam_features, encs_prev, decs_prev = [], [], []
684 |
685 | for idx_stage in range(self.num_stages):
686 | # Input convolution, get multi-scale input features
687 | x_scales = []
688 | for i in range(self.num_supervision_scales):
689 | x_scale = Conv3x3(
690 | (2**i) * self.features,
691 | use_bias=self.use_bias,
692 | name=f"stage_{idx_stage}_input_conv_{i}")(
693 | shortcuts[i])
694 |
695 | # If later stages, fuse input features with SAM features from prev stage
696 | if idx_stage > 0:
697 | # use larger blocksize at high-res stages
698 | if self.use_cross_gating:
699 | block_size = self.block_size_hr if i < self.high_res_stages else self.block_size_lr
700 | grid_size = self.grid_size_hr if i < self.high_res_stages else self.block_size_lr
701 | x_scale, _ = CrossGatingBlock(
702 | features=(2**i) * self.features,
703 | block_size=block_size,
704 | grid_size=grid_size,
705 | dropout_rate=self.dropout_rate,
706 | input_proj_factor=self.input_proj_factor,
707 | upsample_y=False,
708 | use_bias=self.use_bias,
709 | name=f"stage_{idx_stage}_input_fuse_sam_{i}")(
710 | x_scale, sam_features.pop(), deterministic=not train)
711 | else:
712 | x_scale = Conv1x1(
713 | (2**i) * self.features,
714 | use_bias=self.use_bias,
715 | name=f"stage_{idx_stage}_input_catconv_{i}")(
716 | jnp.concatenate(
717 | [x_scale, sam_features.pop()], axis=-1))
718 |
719 | x_scales.append(x_scale)
720 |
721 | # start encoder blocks
722 | encs = []
723 | x = x_scales[0] # First full-scale input feature
724 |
725 | for i in range(self.depth): # 0, 1, 2
726 | # use larger blocksize at high-res stages, vice versa.
727 | block_size = self.block_size_hr if i < self.high_res_stages else self.block_size_lr
728 | grid_size = self.grid_size_hr if i < self.high_res_stages else self.block_size_lr
729 | use_cross_gating_layer = True if idx_stage > 0 else False
730 |
731 | # Multi-scale input if multi-scale supervision
732 | x_scale = x_scales[i] if i < self.num_supervision_scales else None
733 |
734 | # UNet Encoder block
735 | enc_prev = encs_prev.pop() if idx_stage > 0 else None
736 | dec_prev = decs_prev.pop() if idx_stage > 0 else None
737 |
738 | x, bridge = UNetEncoderBlock(
739 | features=(2**i) * self.features,
740 | num_groups=self.num_groups,
741 | downsample=True,
742 | lrelu_slope=self.lrelu_slope,
743 | block_size=block_size,
744 | grid_size=grid_size,
745 | block_gmlp_factor=self.block_gmlp_factor,
746 | grid_gmlp_factor=self.grid_gmlp_factor,
747 | input_proj_factor=self.input_proj_factor,
748 | channels_reduction=self.channels_reduction,
749 | use_global_mlp=self.use_global_mlp,
750 | dropout_rate=self.dropout_rate,
751 | use_bias=self.use_bias,
752 | use_cross_gating=use_cross_gating_layer,
753 | name=f"stage_{idx_stage}_encoder_block_{i}")(
754 | x,
755 | skip=x_scale,
756 | enc=enc_prev,
757 | dec=dec_prev,
758 | deterministic=not train)
759 |
760 | # Cache skip signals
761 | encs.append(bridge)
762 |
763 | # Global MLP bottleneck blocks
764 | for i in range(self.num_bottleneck_blocks):
765 | x = BottleneckBlock(
766 | block_size=self.block_size_lr,
767 | grid_size=self.block_size_lr,
768 | features=(2**(self.depth - 1)) * self.features,
769 | num_groups=self.num_groups,
770 | block_gmlp_factor=self.block_gmlp_factor,
771 | grid_gmlp_factor=self.grid_gmlp_factor,
772 | input_proj_factor=self.input_proj_factor,
773 | dropout_rate=self.dropout_rate,
774 | use_bias=self.use_bias,
775 | channels_reduction=self.channels_reduction,
776 | name=f"stage_{idx_stage}_global_block_{i}")(
777 | x, deterministic=not train)
778 | # cache global feature for cross-gating
779 | global_feature = x
780 |
781 | # start cross gating. Use multi-scale feature fusion
782 | skip_features = []
783 | for i in reversed(range(self.depth)): # 2, 1, 0
784 | # use larger blocksize at high-res stages
785 | block_size = self.block_size_hr if i < self.high_res_stages else self.block_size_lr
786 | grid_size = self.grid_size_hr if i < self.high_res_stages else self.block_size_lr
787 |
788 | # get additional multi-scale signals
789 | signal = jnp.concatenate([
790 | UpSampleRatio(
791 | (2**i) * self.features,
792 | ratio=2**(j - i),
793 | use_bias=self.use_bias)(enc) for j, enc in enumerate(encs)
794 | ],
795 | axis=-1)
796 |
797 | # Use cross-gating to cross modulate features
798 | if self.use_cross_gating:
799 | skips, global_feature = CrossGatingBlock(
800 | features=(2**i) * self.features,
801 | block_size=block_size,
802 | grid_size=grid_size,
803 | input_proj_factor=self.input_proj_factor,
804 | dropout_rate=self.dropout_rate,
805 | upsample_y=True,
806 | use_bias=self.use_bias,
807 | name=f"stage_{idx_stage}_cross_gating_block_{i}")(
808 | signal, global_feature, deterministic=not train)
809 | else:
810 | skips = Conv1x1(
811 | (2**i) * self.features, use_bias=self.use_bias)(
812 | signal)
813 | skips = Conv3x3((2**i) * self.features, use_bias=self.use_bias)(skips)
814 |
815 | skip_features.append(skips)
816 |
817 | # start decoder. Multi-scale feature fusion of cross-gated features
818 | outputs, decs, sam_features = [], [], []
819 | for i in reversed(range(self.depth)):
820 | # use larger blocksize at high-res stages
821 | block_size = self.block_size_hr if i < self.high_res_stages else self.block_size_lr
822 | grid_size = self.grid_size_hr if i < self.high_res_stages else self.block_size_lr
823 |
824 | # get multi-scale skip signals from cross-gating block
825 | signal = jnp.concatenate([
826 | UpSampleRatio(
827 | (2**i) * self.features,
828 | ratio=2**(self.depth - j - 1 - i),
829 | use_bias=self.use_bias)(skip)
830 | for j, skip in enumerate(skip_features)
831 | ],
832 | axis=-1)
833 |
834 | # Decoder block
835 | x = UNetDecoderBlock(
836 | features=(2**i) * self.features,
837 | num_groups=self.num_groups,
838 | lrelu_slope=self.lrelu_slope,
839 | block_size=block_size,
840 | grid_size=grid_size,
841 | block_gmlp_factor=self.block_gmlp_factor,
842 | grid_gmlp_factor=self.grid_gmlp_factor,
843 | input_proj_factor=self.input_proj_factor,
844 | channels_reduction=self.channels_reduction,
845 | use_global_mlp=self.use_global_mlp,
846 | dropout_rate=self.dropout_rate,
847 | use_bias=self.use_bias,
848 | name=f"stage_{idx_stage}_decoder_block_{i}")(
849 | x, bridge=signal, deterministic=not train)
850 |
851 | # Cache decoder features for later-stage's usage
852 | decs.append(x)
853 |
854 | # output conv, if not final stage, use supervised-attention-block.
855 | if i < self.num_supervision_scales:
856 | if idx_stage < self.num_stages - 1: # not last stage, apply SAM
857 | sam, output = SAM(
858 | (2**i) * self.features,
859 | output_channels=self.num_outputs,
860 | use_bias=self.use_bias,
861 | name=f"stage_{idx_stage}_supervised_attention_module_{i}")(
862 | x, shortcuts[i], train=train)
863 | outputs.append(output)
864 | sam_features.append(sam)
865 | else: # Last stage, apply output convolutions
866 | output = Conv3x3(self.num_outputs,
867 | use_bias=self.use_bias,
868 | name=f"stage_{idx_stage}_output_conv_{i}")(x)
869 | output = output + shortcuts[i]
870 | outputs.append(output)
871 | # Cache encoder and decoder features for later-stage's usage
872 | encs_prev = encs[::-1]
873 | decs_prev = decs
874 |
875 | # Store outputs
876 | outputs_all.append(outputs)
877 | return outputs_all
878 |
879 |
880 | def Model(*, variant=None, **kw):
881 | """Factory function to easily create a Model variant like "S".
882 |
883 | Every model file should have this Model() function that returns the flax
884 | model function. The function name should be fixed.
885 |
886 | Args:
887 | variant: UNet model variants. Options: 'S-1' | 'S-2' | 'S-3'
888 | | 'M-1' | 'M-2' | 'M-3'
889 | **kw: Other UNet config dicts.
890 |
891 | Returns:
892 | The MAXIM() model function
893 | """
894 |
895 | if variant is not None:
896 | config = {
897 | # params: 6.108515000000001 M, GFLOPS: 93.163716608
898 | "S-1": {
899 | "features": 32,
900 | "depth": 3,
901 | "num_stages": 1,
902 | "num_groups": 2,
903 | "num_bottleneck_blocks": 2,
904 | "block_gmlp_factor": 2,
905 | "grid_gmlp_factor": 2,
906 | "input_proj_factor": 2,
907 | "channels_reduction": 4,
908 | },
909 | # params: 13.35383 M, GFLOPS: 206.743273472
910 | "S-2": {
911 | "features": 32,
912 | "depth": 3,
913 | "num_stages": 2,
914 | "num_groups": 2,
915 | "num_bottleneck_blocks": 2,
916 | "block_gmlp_factor": 2,
917 | "grid_gmlp_factor": 2,
918 | "input_proj_factor": 2,
919 | "channels_reduction": 4,
920 | },
921 | # params: 20.599145 M, GFLOPS: 320.32194560000005
922 | "S-3": {
923 | "features": 32,
924 | "depth": 3,
925 | "num_stages": 3,
926 | "num_groups": 2,
927 | "num_bottleneck_blocks": 2,
928 | "block_gmlp_factor": 2,
929 | "grid_gmlp_factor": 2,
930 | "input_proj_factor": 2,
931 | "channels_reduction": 4,
932 | },
933 | # params: 19.361219000000002 M, 308.495712256 GFLOPs
934 | "M-1": {
935 | "features": 64,
936 | "depth": 3,
937 | "num_stages": 1,
938 | "num_groups": 2,
939 | "num_bottleneck_blocks": 2,
940 | "block_gmlp_factor": 2,
941 | "grid_gmlp_factor": 2,
942 | "input_proj_factor": 2,
943 | "channels_reduction": 4,
944 | },
945 | # params: 40.83911 M, 675.25541888 GFLOPs
946 | "M-2": {
947 | "features": 64,
948 | "depth": 3,
949 | "num_stages": 2,
950 | "num_groups": 2,
951 | "num_bottleneck_blocks": 2,
952 | "block_gmlp_factor": 2,
953 | "grid_gmlp_factor": 2,
954 | "input_proj_factor": 2,
955 | "channels_reduction": 4,
956 | },
957 | # params: 62.317001 M, 1042.014666752 GFLOPs
958 | "M-3": {
959 | "features": 64,
960 | "depth": 3,
961 | "num_stages": 3,
962 | "num_groups": 2,
963 | "num_bottleneck_blocks": 2,
964 | "block_gmlp_factor": 2,
965 | "grid_gmlp_factor": 2,
966 | "input_proj_factor": 2,
967 | "channels_reduction": 4,
968 | },
969 | }[variant]
970 |
971 | for k, v in config.items():
972 | kw.setdefault(k, v)
973 |
974 | return MAXIM(**kw)
975 |
--------------------------------------------------------------------------------
/maxim/predict.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | import os
16 | os.environ["XLA_FLAGS"] = "--xla_gpu_force_compilation_parallelism=1"
17 |
18 | import subprocess
19 | subprocess.call(["pip", "install", "."])
20 |
21 | import numpy as np
22 | from PIL import Image
23 | import importlib
24 | import ml_collections
25 | import tempfile
26 | import jax.numpy as jnp
27 | import flax
28 | from cog import BasePredictor, Path, Input, BaseModel
29 |
30 | from maxim.run_eval import (
31 | _MODEL_FILENAME,
32 | _MODEL_VARIANT_DICT,
33 | _MODEL_CONFIGS,
34 | get_params,
35 | mod_padding_symmetric,
36 | make_shape_even,
37 | augment_image,
38 | )
39 |
40 |
41 | class Predictor(BasePredictor):
42 | def setup(self):
43 |
44 | self.params = {
45 | "Image Denoising": get_params("checkpoints/denoising-SIDD/checkpoint.npz"),
46 | "Image Deblurring (GoPro)": get_params(
47 | "checkpoints/debluring-GoPro/checkpoint.npz"
48 | ),
49 | "Image Deblurring (REDS)": get_params(
50 | "checkpoints/debluring-REDS/checkpoint.npz"
51 | ),
52 | "Image Deblurring (RealBlur_R)": get_params(
53 | "checkpoints/debluring-Real-Blur-R/checkpoint.npz"
54 | ),
55 | "Image Deblurring (RealBlur_J)": get_params(
56 | "checkpoints/debluring-Real-Blur-J/checkpoint.npz"
57 | ),
58 | "Image Deraining (Rain streak)": get_params(
59 | "checkpoints/deraining-Rain13k/checkpoint.npz"
60 | ),
61 | "Image Deraining (Rain drop)": get_params(
62 | "checkpoints/deraining-Raindrop/checkpoint.npz"
63 | ),
64 | "Image Dehazing (Indoor)": get_params(
65 | "checkpoints/dehazing-RESIDE-Indoor/checkpoint.npz"
66 | ),
67 | "Image Dehazing (Outdoor)": get_params(
68 | "checkpoints/dehazing-RESIDE-Outdoor/checkpoint.npz"
69 | ),
70 | "Image Enhancement (Low-light)": get_params(
71 | "checkpoints/enhancement-LOL/checkpoint.npz"
72 | ),
73 | "Image Enhancement (Retouching)": get_params(
74 | "checkpoints/enhancement-FiveK/checkpoint.npz"
75 | ),
76 | }
77 |
78 | model_mod = importlib.import_module(f"maxim.models.{_MODEL_FILENAME}")
79 | self.models = {}
80 | for task in _MODEL_VARIANT_DICT.keys():
81 | model_configs = ml_collections.ConfigDict(_MODEL_CONFIGS)
82 | model_configs.variant = _MODEL_VARIANT_DICT[task]
83 | self.models[task] = model_mod.Model(**model_configs)
84 |
85 | def predict(
86 | self,
87 | model: str = Input(
88 | choices=[
89 | "Image Denoising",
90 | "Image Deblurring (GoPro)",
91 | "Image Deblurring (REDS)",
92 | "Image Deblurring (RealBlur_R)",
93 | "Image Deblurring (RealBlur_J)",
94 | "Image Deraining (Rain streak)",
95 | "Image Deraining (Rain drop)",
96 | "Image Dehazing (Indoor)",
97 | "Image Dehazing (Outdoor)",
98 | "Image Enhancement (Low-light)",
99 | "Image Enhancement (Retouching)",
100 | ],
101 | description="Choose a model.",
102 | ),
103 | image: Path = Input(
104 | description="Input image.",
105 | ),
106 | ) -> Path:
107 |
108 | params = self.params[model]
109 | task = model.split()[1]
110 | model = self.models[task]
111 |
112 | input_img = (
113 | np.asarray(Image.open(str(image)).convert("RGB"), np.float32) / 255.0
114 | )
115 |
116 | # Padding images to have even shapes
117 | height, width = input_img.shape[0], input_img.shape[1]
118 | input_img = make_shape_even(input_img)
119 | height_even, width_even = input_img.shape[0], input_img.shape[1]
120 |
121 | # padding images to be multiplies of 64
122 | input_img = mod_padding_symmetric(input_img, factor=64)
123 | input_img = np.expand_dims(input_img, axis=0)
124 |
125 | # handle multi-stage outputs, obtain the last scale output of last stage
126 | preds = model.apply({"params": flax.core.freeze(params)}, input_img)
127 | if isinstance(preds, list):
128 | preds = preds[-1]
129 | if isinstance(preds, list):
130 | preds = preds[-1]
131 |
132 | preds = np.array(preds[0], np.float32)
133 |
134 | # unpad images to get the original resolution
135 | new_height, new_width = preds.shape[0], preds.shape[1]
136 | h_start = new_height // 2 - height_even // 2
137 | h_end = h_start + height
138 | w_start = new_width // 2 - width_even // 2
139 | w_end = w_start + width
140 | preds = preds[h_start:h_end, w_start:w_end, :]
141 |
142 | # save files
143 | out_path = Path(tempfile.mkdtemp()) / "output.png"
144 | Image.fromarray(
145 | np.array((np.clip(preds, 0.0, 1.0) * 255.0).astype(jnp.uint8))
146 | ).save(str(out_path))
147 |
148 | return out_path
--------------------------------------------------------------------------------
/maxim/run_eval.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | """Run evaluation."""
16 |
17 | import collections
18 | import importlib
19 | import io
20 | import os
21 |
22 | from absl import app
23 | from absl import flags
24 | import flax
25 | import jax.numpy as jnp
26 | import ml_collections
27 | import numpy as np
28 | from PIL import Image
29 | import tensorflow as tf
30 |
31 | FLAGS = flags.FLAGS
32 |
33 | flags.DEFINE_enum(
34 | 'task', 'Denoising',
35 | ['Denoising', 'Deblurring', 'Deraining', 'Dehazing', 'Enhancement'],
36 | 'Task to run.')
37 | flags.DEFINE_string('ckpt_path', '', 'Path to checkpoint.')
38 | flags.DEFINE_string('input_dir', '', 'Input dir to the test set.')
39 | flags.DEFINE_string('output_dir', '', 'Output dir to store predicted images.')
40 | flags.DEFINE_boolean('has_target', True, 'Whether has corresponding gt image.')
41 | flags.DEFINE_boolean('save_images', True, 'Dump predicted images.')
42 | flags.DEFINE_boolean('geometric_ensemble', False,
43 | 'Whether use ensemble infernce.')
44 |
45 | _MODEL_FILENAME = 'maxim'
46 |
47 | _MODEL_VARIANT_DICT = {
48 | 'Denoising': 'S-3',
49 | 'Deblurring': 'S-3',
50 | 'Deraining': 'S-2',
51 | 'Dehazing': 'S-2',
52 | 'Enhancement': 'S-2',
53 | }
54 |
55 | _MODEL_CONFIGS = {
56 | 'variant': '',
57 | 'dropout_rate': 0.0,
58 | 'num_outputs': 3,
59 | 'use_bias': True,
60 | 'num_supervision_scales': 3,
61 | }
62 |
63 |
64 | def recover_tree(keys, values):
65 | """Recovers a tree as a nested dict from flat names and values.
66 |
67 | This function is useful to analyze checkpoints that are saved by our programs
68 | without need to access the exact source code of the experiment. In particular,
69 | it can be used to extract an reuse various subtrees of the scheckpoint, e.g.
70 | subtree of parameters.
71 | Args:
72 | keys: a list of keys, where '/' is used as separator between nodes.
73 | values: a list of leaf values.
74 | Returns:
75 | A nested tree-like dict.
76 | """
77 | tree = {}
78 | sub_trees = collections.defaultdict(list)
79 | for k, v in zip(keys, values):
80 | if '/' not in k:
81 | tree[k] = v
82 | else:
83 | k_left, k_right = k.split('/', 1)
84 | sub_trees[k_left].append((k_right, v))
85 | for k, kv_pairs in sub_trees.items():
86 | k_subtree, v_subtree = zip(*kv_pairs)
87 | tree[k] = recover_tree(k_subtree, v_subtree)
88 | return tree
89 |
90 |
91 | def mod_padding_symmetric(image, factor=64):
92 | """Padding the image to be divided by factor."""
93 | height, width = image.shape[0], image.shape[1]
94 | height_pad, width_pad = ((height + factor) // factor) * factor, (
95 | (width + factor) // factor) * factor
96 | padh = height_pad - height if height % factor != 0 else 0
97 | padw = width_pad - width if width % factor != 0 else 0
98 | image = jnp.pad(
99 | image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)],
100 | mode='reflect')
101 | return image
102 |
103 |
104 | def get_params(ckpt_path):
105 | """Get params checkpoint."""
106 |
107 | with tf.io.gfile.GFile(ckpt_path, 'rb') as f:
108 | data = f.read()
109 | values = np.load(io.BytesIO(data))
110 | params = recover_tree(*zip(*values.items()))
111 | params = params['opt']['target']
112 |
113 | return params
114 |
115 |
116 | def calculate_psnr(img1, img2, crop_border, test_y_channel=False):
117 | """Calculate PSNR (Peak Signal-to-Noise Ratio).
118 |
119 | Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
120 | Args:
121 | img1 (ndarray): Images with range [0, 255].
122 | img2 (ndarray): Images with range [0, 255].
123 | crop_border (int): Cropped pixels in each edge of an image. These
124 | pixels are not involved in the PSNR calculation.
125 | test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
126 | Returns:
127 | float: psnr result.
128 | """
129 | assert img1.shape == img2.shape, (
130 | f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
131 | img1 = img1.astype(np.float64)
132 | img2 = img2.astype(np.float64)
133 |
134 | if crop_border != 0:
135 | img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...]
136 | img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
137 |
138 | if test_y_channel:
139 | img1 = to_y_channel(img1)
140 | img2 = to_y_channel(img2)
141 |
142 | mse = np.mean((img1 - img2)**2)
143 | if mse == 0:
144 | return float('inf')
145 | return 20. * np.log10(255. / np.sqrt(mse))
146 |
147 |
148 | def _convert_input_type_range(img):
149 | """Convert the type and range of the input image.
150 |
151 | It converts the input image to np.float32 type and range of [0, 1].
152 | It is mainly used for pre-processing the input image in colorspace
153 | convertion functions such as rgb2ycbcr and ycbcr2rgb.
154 | Args:
155 | img (ndarray): The input image. It accepts:
156 | 1. np.uint8 type with range [0, 255];
157 | 2. np.float32 type with range [0, 1].
158 | Returns:
159 | (ndarray): The converted image with type of np.float32 and range of
160 | [0, 1].
161 | """
162 | img_type = img.dtype
163 | img = img.astype(np.float32)
164 | if img_type == np.float32:
165 | pass
166 | elif img_type == np.uint8:
167 | img /= 255.
168 | else:
169 | raise TypeError('The img type should be np.float32 or np.uint8, '
170 | f'but got {img_type}')
171 | return img
172 |
173 |
174 | def _convert_output_type_range(img, dst_type):
175 | """Convert the type and range of the image according to dst_type.
176 |
177 | It converts the image to desired type and range. If `dst_type` is np.uint8,
178 | images will be converted to np.uint8 type with range [0, 255]. If
179 | `dst_type` is np.float32, it converts the image to np.float32 type with
180 | range [0, 1].
181 | It is mainly used for post-processing images in colorspace convertion
182 | functions such as rgb2ycbcr and ycbcr2rgb.
183 | Args:
184 | img (ndarray): The image to be converted with np.float32 type and
185 | range [0, 255].
186 | dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
187 | converts the image to np.uint8 type with range [0, 255]. If
188 | dst_type is np.float32, it converts the image to np.float32 type
189 | with range [0, 1].
190 | Returns:
191 | (ndarray): The converted image with desired type and range.
192 | """
193 | if dst_type not in (np.uint8, np.float32):
194 | raise TypeError('The dst_type should be np.float32 or np.uint8, '
195 | f'but got {dst_type}')
196 | if dst_type == np.uint8:
197 | img = img.round()
198 | else:
199 | img /= 255.
200 |
201 | return img.astype(dst_type)
202 |
203 |
204 | def rgb2ycbcr(img, y_only=False):
205 | """Convert a RGB image to YCbCr image.
206 |
207 | This function produces the same results as Matlab's `rgb2ycbcr` function.
208 | It implements the ITU-R BT.601 conversion for standard-definition
209 | television. See more details in
210 | https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
211 | It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
212 | In OpenCV, it implements a JPEG conversion. See more details in
213 | https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
214 |
215 | Args:
216 | img (ndarray): The input image. It accepts:
217 | 1. np.uint8 type with range [0, 255];
218 | 2. np.float32 type with range [0, 1].
219 | y_only (bool): Whether to only return Y channel. Default: False.
220 | Returns:
221 | ndarray: The converted YCbCr image. The output image has the same type
222 | and range as input image.
223 | """
224 | img_type = img.dtype
225 | img = _convert_input_type_range(img)
226 | if y_only:
227 | out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
228 | else:
229 | out_img = np.matmul(img,
230 | [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
231 | [24.966, 112.0, -18.214]]) + [16, 128, 128]
232 | out_img = _convert_output_type_range(out_img, img_type)
233 | return out_img
234 |
235 |
236 | def to_y_channel(img):
237 | """Change to Y channel of YCbCr.
238 |
239 | Args:
240 | img (ndarray): Images with range [0, 255].
241 | Returns:
242 | (ndarray): Images with range [0, 255] (float type) without round.
243 | """
244 | img = img.astype(np.float32) / 255.
245 | if img.ndim == 3 and img.shape[2] == 3:
246 | img = rgb2ycbcr(img, y_only=True)
247 | img = img[..., None]
248 | return img * 255.
249 |
250 |
251 | def augment_image(image, times=8):
252 | """Geometric augmentation."""
253 | if times == 4: # only rotate image
254 | images = []
255 | for k in range(0, 4):
256 | images.append(np.rot90(image, k=k))
257 | images = np.stack(images, axis=0)
258 | elif times == 8: # roate and flip image
259 | images = []
260 | for k in range(0, 4):
261 | images.append(np.rot90(image, k=k))
262 | image = np.fliplr(image)
263 | for k in range(0, 4):
264 | images.append(np.rot90(image, k=k))
265 | images = np.stack(images, axis=0)
266 | else:
267 | raise Exception(f'Error times: {times}')
268 | return images
269 |
270 |
271 | def deaugment_image(images, times=8):
272 | """Reverse the geometric augmentation."""
273 |
274 | if times == 4: # only rotate image
275 | image = []
276 | for k in range(0, 4):
277 | image.append(np.rot90(images[k], k=4-k))
278 | image = np.stack(image, axis=0)
279 | image = np.mean(image, axis=0)
280 | elif times == 8: # roate and flip image
281 | image = []
282 | for k in range(0, 4):
283 | image.append(np.rot90(images[k], k=4-k))
284 | for k in range(0, 4):
285 | image.append(np.fliplr(np.rot90(images[4+k], k=4-k)))
286 | image = np.mean(image, axis=0)
287 | else:
288 | raise Exception(f'Error times: {times}')
289 | return image
290 |
291 |
292 | def is_image_file(filename):
293 | """Check if it is an valid image file by extension."""
294 | return any(
295 | filename.endswith(extension)
296 | for extension in ['jpeg', 'JPEG', 'jpg', 'png', 'JPG', 'PNG', 'gif'])
297 |
298 |
299 | def save_img(img, pth):
300 | """Save an image to disk.
301 |
302 | Args:
303 | img: jnp.ndarry, [height, width, channels], img will be clipped to [0, 1]
304 | before saved to pth.
305 | pth: string, path to save the image to.
306 | """
307 | Image.fromarray(np.array(
308 | (np.clip(img, 0., 1.) * 255.).astype(jnp.uint8))).save(pth, 'PNG')
309 |
310 |
311 | def make_shape_even(image):
312 | """Pad the image to have even shapes."""
313 | height, width = image.shape[0], image.shape[1]
314 | padh = 1 if height % 2 != 0 else 0
315 | padw = 1 if width % 2 != 0 else 0
316 | image = jnp.pad(image, [(0, padh), (0, padw), (0, 0)], mode='reflect')
317 | return image
318 |
319 |
320 | def main(_):
321 | params = get_params(FLAGS.ckpt_path)
322 |
323 | if FLAGS.save_images:
324 | os.makedirs(FLAGS.output_dir, exist_ok=True)
325 |
326 | # sorted is important for continuning an inference job.
327 | filepath = sorted(os.listdir(os.path.join(FLAGS.input_dir, 'input')))
328 | input_filenames = [
329 | os.path.join(FLAGS.input_dir, 'input', x)
330 | for x in filepath
331 | if is_image_file(x)
332 | ]
333 | if FLAGS.has_target:
334 | target_filenames = [
335 | os.path.join(FLAGS.input_dir, 'target', x)
336 | for x in filepath
337 | if is_image_file(x)
338 | ]
339 | num_images = len(input_filenames)
340 |
341 | model_mod = importlib.import_module(f'maxim.models.{_MODEL_FILENAME}')
342 | model_configs = ml_collections.ConfigDict(_MODEL_CONFIGS)
343 | model_configs.variant = _MODEL_VARIANT_DICT[FLAGS.task]
344 | model = model_mod.Model(**model_configs)
345 |
346 | psnr_all = []
347 |
348 | def _process_file(i):
349 | print(f'Processing {i + 1} / {num_images}...')
350 | input_file = input_filenames[i]
351 | input_img = np.asarray(Image.open(input_file).convert('RGB'),
352 | np.float32) / 255.
353 | if FLAGS.has_target:
354 | target_file = target_filenames[i]
355 | target_img = np.asarray(Image.open(target_file).convert('RGB'),
356 | np.float32) / 255.
357 |
358 | # Padding images to have even shapes
359 | height, width = input_img.shape[0], input_img.shape[1]
360 | input_img = make_shape_even(input_img)
361 | height_even, width_even = input_img.shape[0], input_img.shape[1]
362 |
363 | # padding images to be multiplies of 64
364 | input_img = mod_padding_symmetric(input_img, factor=64)
365 |
366 | if FLAGS.geometric_ensemble:
367 | input_img = augment_image(input_img, FLAGS.ensemble_times)
368 | else:
369 | input_img = np.expand_dims(input_img, axis=0)
370 |
371 | # handle multi-stage outputs, obtain the last scale output of last stage
372 | preds = model.apply({'params': flax.core.freeze(params)}, input_img)
373 | if isinstance(preds, list):
374 | preds = preds[-1]
375 | if isinstance(preds, list):
376 | preds = preds[-1]
377 |
378 | # De-ensemble by averaging inferenced results.
379 | if FLAGS.geometric_ensemble:
380 | preds = deaugment_image(preds, FLAGS.ensemble_times)
381 | else:
382 | preds = np.array(preds[0], np.float32)
383 |
384 | # unpad images to get the original resolution
385 | new_height, new_width = preds.shape[0], preds.shape[1]
386 | h_start = new_height // 2 - height_even // 2
387 | h_end = h_start + height
388 | w_start = new_width // 2 - width_even // 2
389 | w_end = w_start + width
390 | preds = preds[h_start:h_end, w_start:w_end, :]
391 |
392 | # print PSNR scores
393 | if FLAGS.has_target:
394 | psnr = calculate_psnr(
395 | target_img * 255., preds * 255., crop_border=0, test_y_channel=False)
396 | print(f'{i}th image: psnr = {psnr:.4f}')
397 | else:
398 | psnr = -1
399 |
400 | # save files
401 | basename = os.path.basename(input_file)
402 | if FLAGS.save_images:
403 | save_pth = os.path.join(FLAGS.output_dir, basename)
404 | save_img(preds, save_pth)
405 |
406 | return psnr
407 |
408 | for i in range(num_images):
409 | psnr = _process_file(i)
410 | psnr_all.append(psnr)
411 |
412 | psnr_all = np.asarray(psnr_all)
413 |
414 | print(f'average psnr = {np.sum(psnr_all)/num_images:.4f}')
415 | print(f'std psnr = {np.std(psnr_all):.4f}')
416 |
417 |
418 | if __name__ == '__main__':
419 | app.run(main)
420 |
--------------------------------------------------------------------------------
/maxim/test_maxim.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | """Test maxim."""
16 | from absl.testing import absltest
17 |
18 | if __name__ == '__main__':
19 | absltest.main()
20 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | absl-py>=0.12.0
2 | chex>=0.0.7
3 | clu>=0.0.3
4 | einops>=0.3.0
5 | flax>=0.3.3
6 | ml-collections==0.1.0
7 | numpy>=1.19.5
8 | pandas>=1.1.0
9 | tensorflow>=2.0.0-beta1
10 | jax>=0.1.55
11 | jaxlib>=0.1.37
12 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 Google LLC.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | """Install MAXIM."""
16 |
17 | import setuptools
18 |
19 | # Get install requirements from the REQUIREMENTS file.
20 | with open('requirements.txt') as fp:
21 | _REQUIREMENTS = fp.read().splitlines()
22 |
23 | # Get the long description from the README file.
24 | with open('README.md') as fp:
25 | _LONG_DESCRIPTION = fp.read()
26 |
27 | setuptools.setup(
28 | name='maxim',
29 | version='1.0.0',
30 | description='MAXIM: Multi-Axis MLP for Image Processing',
31 | long_description=_LONG_DESCRIPTION,
32 | long_description_content_type='text/markdown',
33 | author='Google Inc.',
34 | author_email='no-reply@google.com',
35 | url='http://github.com/google-research/maxim',
36 | license='Apache 2.0',
37 | packages=[
38 | 'maxim', 'maxim.models',
39 | ],
40 | scripts=[],
41 | install_requires=_REQUIREMENTS,
42 | keywords='deeplearning machinelearning transformer mlp lowlevel imageprocessing',
43 | )
44 |
--------------------------------------------------------------------------------
Image Enhancement (click to expand)
219 | 220 |
221 |