├── maxim ├── __init__.py ├── tests │ ├── __init__.py │ ├── test_maxim_variable_img_length.py │ ├── conftest.py │ └── test_block_operations.py ├── blocks │ ├── __init__.py │ ├── others.py │ ├── bottleneck.py │ ├── block_gating.py │ ├── grid_gating.py │ ├── unet.py │ ├── attentions.py │ └── misc_gating.py ├── configs.py ├── layers.py └── maxim.py ├── requirements.txt ├── images ├── maxim.gif ├── overview.png ├── Deraining │ └── input │ │ ├── 0.jpg │ │ ├── 1.png │ │ ├── 15.png │ │ └── 55.png ├── Enhancement │ └── input │ │ ├── 1.png │ │ ├── 111.png │ │ ├── 748.png │ │ └── a4541-DSC_0040-2.png ├── Dehazing │ └── input │ │ ├── 1440_10.png │ │ ├── 1444_10.png │ │ ├── 0003_0.8_0.2.png │ │ ├── 0014_0.8_0.12.png │ │ ├── 0048_0.9_0.2.png │ │ └── 0010_0.95_0.16.png ├── Denoising │ └── input │ │ ├── 0003_30.png │ │ ├── 0011_23.png │ │ ├── 0013_19.png │ │ └── 0039_04.png └── Deblurring │ └── input │ ├── 1fromGOPR0950.png │ ├── 109fromGOPR1096.MP4.png │ ├── 110fromGOPR1087.MP4.png │ └── 1fromGOPR1096.MP4.png ├── hub_utilities ├── obtain_sm.py ├── export_for_hub.py └── generate_doc.py ├── create_maxim_model.py ├── benchmark_xla.py ├── convert_to_tf.py ├── README.md ├── LICENSE ├── notebooks ├── inference.ipynb └── inference-dynamic-resize.ipynb └── run_eval.py /maxim/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /maxim/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /maxim/blocks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow==2.10.0 2 | einops==0.5.0 3 | huggingface_hub -------------------------------------------------------------------------------- /images/maxim.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/maxim.gif -------------------------------------------------------------------------------- /images/overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/overview.png -------------------------------------------------------------------------------- /images/Deraining/input/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deraining/input/0.jpg -------------------------------------------------------------------------------- /images/Deraining/input/1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deraining/input/1.png -------------------------------------------------------------------------------- /images/Deraining/input/15.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deraining/input/15.png -------------------------------------------------------------------------------- /images/Deraining/input/55.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deraining/input/55.png -------------------------------------------------------------------------------- /images/Enhancement/input/1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Enhancement/input/1.png -------------------------------------------------------------------------------- /images/Enhancement/input/111.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Enhancement/input/111.png -------------------------------------------------------------------------------- /images/Enhancement/input/748.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Enhancement/input/748.png -------------------------------------------------------------------------------- /images/Dehazing/input/1440_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/1440_10.png -------------------------------------------------------------------------------- /images/Dehazing/input/1444_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/1444_10.png -------------------------------------------------------------------------------- /images/Denoising/input/0003_30.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Denoising/input/0003_30.png -------------------------------------------------------------------------------- /images/Denoising/input/0011_23.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Denoising/input/0011_23.png -------------------------------------------------------------------------------- /images/Denoising/input/0013_19.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Denoising/input/0013_19.png -------------------------------------------------------------------------------- /images/Denoising/input/0039_04.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Denoising/input/0039_04.png -------------------------------------------------------------------------------- /images/Dehazing/input/0003_0.8_0.2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/0003_0.8_0.2.png -------------------------------------------------------------------------------- /images/Dehazing/input/0014_0.8_0.12.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/0014_0.8_0.12.png -------------------------------------------------------------------------------- /images/Dehazing/input/0048_0.9_0.2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/0048_0.9_0.2.png -------------------------------------------------------------------------------- /images/Deblurring/input/1fromGOPR0950.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deblurring/input/1fromGOPR0950.png -------------------------------------------------------------------------------- /images/Dehazing/input/0010_0.95_0.16.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Dehazing/input/0010_0.95_0.16.png -------------------------------------------------------------------------------- /images/Deblurring/input/109fromGOPR1096.MP4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deblurring/input/109fromGOPR1096.MP4.png -------------------------------------------------------------------------------- /images/Deblurring/input/110fromGOPR1087.MP4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deblurring/input/110fromGOPR1087.MP4.png -------------------------------------------------------------------------------- /images/Deblurring/input/1fromGOPR1096.MP4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Deblurring/input/1fromGOPR1096.MP4.png -------------------------------------------------------------------------------- /images/Enhancement/input/a4541-DSC_0040-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sayakpaul/maxim-tf/main/images/Enhancement/input/a4541-DSC_0040-2.png -------------------------------------------------------------------------------- /maxim/tests/test_maxim_variable_img_length.py: -------------------------------------------------------------------------------- 1 | def test_maxim_variable_length(none_model, random_image_multiple_of_64): 2 | # this line will run several times with random images of different size 3 | # The none_model is instantiated only once, since the fixture scope is session. 4 | out = none_model(random_image_multiple_of_64) 5 | -------------------------------------------------------------------------------- /hub_utilities/obtain_sm.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | GCS_ROOT = "gs://gresearch/maxim/ckpt" 4 | 5 | # From https://github.com/google-research/maxim#results-and-pre-trained-models 6 | DS_TASKS_MAP = { 7 | "Denoising": ["SIDD"], 8 | "Deblurring": [ 9 | "GoPro", 10 | "REDS", 11 | "RealBlur_R", 12 | "RealBlur_J", 13 | ], 14 | "Deraining": ["Rain13k", "Raindrop"], 15 | "Dehazing": [ 16 | "SOTS-Indoor", 17 | "SOTS-Outdoor", 18 | ], 19 | "Enhancement": ["LOL", "FiveK"], 20 | } 21 | 22 | 23 | def main(): 24 | for task in DS_TASKS_MAP: 25 | datasets = DS_TASKS_MAP[task] 26 | 27 | for dataset in datasets: 28 | command = f"python ../convert_to_tf.py -t {task} -c {GCS_ROOT}/{task}/{dataset}/checkpoint.npz" 29 | print(f"Converting for task: {task} and dataset: {dataset}.") 30 | os.system(command) 31 | 32 | 33 | if __name__ == "__main__": 34 | main() 35 | -------------------------------------------------------------------------------- /create_maxim_model.py: -------------------------------------------------------------------------------- 1 | from tensorflow import keras 2 | 3 | from maxim import maxim 4 | from maxim.configs import MAXIM_CONFIGS 5 | 6 | 7 | def Model(variant=None, input_resolution=(None, None), **kw) -> keras.Model: 8 | """Factory function to easily create a Model variant like "S". 9 | 10 | Args: 11 | variant: UNet model variants. Options: 'S-1' | 'S-2' | 'S-3' 12 | | 'M-1' | 'M-2' | 'M-3' 13 | input_resolution: Size of the input images. 14 | **kw: Other UNet config dicts. 15 | 16 | Returns: 17 | The MAXIM model. 18 | """ 19 | 20 | if variant is not None: 21 | config = MAXIM_CONFIGS[variant] 22 | for k, v in config.items(): 23 | kw.setdefault(k, v) 24 | 25 | if "variant" in kw: 26 | _ = kw.pop("variant") 27 | if "input_resolution" in kw: 28 | _ = kw.pop("input_resolution") 29 | model_name = kw.pop("name") 30 | 31 | maxim_model = maxim.MAXIM(**kw) 32 | 33 | inputs = keras.Input((*input_resolution, 3)) 34 | outputs = maxim_model(inputs) 35 | final_model = keras.Model(inputs, outputs, name=f"{model_name}_model") 36 | 37 | return final_model 38 | -------------------------------------------------------------------------------- /hub_utilities/export_for_hub.py: -------------------------------------------------------------------------------- 1 | """Generates .tar.gz archives from SavedModels and serializes them.""" 2 | 3 | 4 | import os 5 | from typing import List 6 | 7 | import tensorflow as tf 8 | 9 | TF_MODEL_ROOT = "gs://maxim-tf" 10 | TAR_ARCHIVES = os.path.join(TF_MODEL_ROOT, "tars/") 11 | 12 | 13 | def prepare_archive(model_name: str) -> None: 14 | """Prepares a tar archive.""" 15 | archive_name = f"{model_name}.tar.gz" 16 | print(f"Archiving to {archive_name}.") 17 | archive_command = f"cd {model_name} && tar -czvf ../{archive_name} *" 18 | os.system(archive_command) 19 | os.system(f"rm -rf {model_name}") 20 | 21 | 22 | def save_to_gcs(model_paths: List[str]) -> None: 23 | """Prepares tar archives and saves them inside a GCS bucket.""" 24 | for path in model_paths: 25 | print(f"Preparing model: {path}.") 26 | model_name = path.strip("/") 27 | abs_model_path = os.path.join(TF_MODEL_ROOT, model_name) 28 | 29 | print(f"Copying from {abs_model_path}.") 30 | os.system(f"gsutil cp -r {abs_model_path} .") 31 | prepare_archive(model_name) 32 | 33 | os.system(f"gsutil -m cp -r *.tar.gz {TAR_ARCHIVES}") 34 | os.system("rm -rf *.tar.gz") 35 | 36 | 37 | model_paths = tf.io.gfile.listdir(TF_MODEL_ROOT) 38 | print(f"Total models: {len(model_paths)}.") 39 | 40 | print("Preparing archives for the classification and feature extractor models.") 41 | save_to_gcs(model_paths) 42 | tar_paths = tf.io.gfile.listdir(TAR_ARCHIVES) 43 | print(f"Total tars: {len(tar_paths)}.") 44 | -------------------------------------------------------------------------------- /benchmark_xla.py: -------------------------------------------------------------------------------- 1 | """ 2 | Script to benchmark the regular MAXIM model in TF and its JiT-compiled variant. 3 | 4 | Expected outputs (benchmarked on my Mac locally): 5 | 6 | ``` 7 | Benchmarking TF model... 8 | Average latency (seconds): 3.1694554823999987. 9 | Benchmarking Jit-compiled TF model... 10 | Average latency (seconds): 1.2475706969000029. 11 | ``` 12 | """ 13 | 14 | 15 | import timeit 16 | 17 | import numpy as np 18 | import tensorflow as tf 19 | 20 | from create_maxim_model import Model 21 | 22 | INPUT_RESOLUTION = 256 23 | 24 | MAXIM_S1 = Model("S-1") 25 | DUMMY_INPUTS = tf.ones((1, INPUT_RESOLUTION, INPUT_RESOLUTION, 3)) 26 | 27 | 28 | def benchmark_regular_model(): 29 | # Warmup 30 | print("Benchmarking TF model...") 31 | for _ in range(2): 32 | _ = MAXIM_S1(DUMMY_INPUTS, training=False) 33 | 34 | # Timing 35 | tf_runtimes = timeit.repeat( 36 | lambda: MAXIM_S1(DUMMY_INPUTS, training=False), number=1, repeat=10 37 | ) 38 | print(f"Average latency (seconds): {np.mean(tf_runtimes)}.") 39 | 40 | 41 | @tf.function(jit_compile=True) 42 | def infer(): 43 | return MAXIM_S1(DUMMY_INPUTS, training=False) 44 | 45 | 46 | def benchmark_xla_model(): 47 | # Warmup 48 | print("Benchmarking Jit-compiled TF model...") 49 | for _ in range(2): 50 | _ = infer() 51 | 52 | # Timing 53 | tf_runtimes = timeit.repeat(lambda: infer(), number=1, repeat=10) 54 | print(f"Average latency (seconds): {np.mean(tf_runtimes)}.") 55 | 56 | 57 | if __name__ == "__main__": 58 | benchmark_regular_model() 59 | benchmark_xla_model() 60 | -------------------------------------------------------------------------------- /maxim/blocks/others.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | import tensorflow as tf 8 | from tensorflow.keras import backend as K 9 | from tensorflow.keras import layers 10 | 11 | from ..layers import Resizing 12 | 13 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 14 | 15 | 16 | def MlpBlock( 17 | mlp_dim: int, 18 | dropout_rate: float = 0.0, 19 | use_bias: bool = True, 20 | name: str = "mlp_block", 21 | ): 22 | """A 1-hidden-layer MLP block, applied over the last dimension.""" 23 | 24 | def apply(x): 25 | d = K.int_shape(x)[-1] 26 | x = layers.Dense(mlp_dim, use_bias=use_bias, name=f"{name}_Dense_0")(x) 27 | x = tf.nn.gelu(x, approximate=True) 28 | x = layers.Dropout(dropout_rate)(x) 29 | x = layers.Dense(d, use_bias=use_bias, name=f"{name}_Dense_1")(x) 30 | return x 31 | 32 | return apply 33 | 34 | 35 | def UpSampleRatio( 36 | num_channels: int, ratio: float, use_bias: bool = True, name: str = "upsample" 37 | ): 38 | """Upsample features given a ratio > 0.""" 39 | 40 | def apply(x): 41 | # Following `jax.image.resize()` 42 | x = Resizing( 43 | ratio=1 / ratio, 44 | method="bilinear", 45 | antialias=True, 46 | name=f"{name}_resizing_{K.get_uid('Resizing')}", 47 | )(x) 48 | 49 | x = Conv1x1(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_0")(x) 50 | return x 51 | 52 | return apply 53 | -------------------------------------------------------------------------------- /maxim/blocks/bottleneck.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | from tensorflow.keras import layers 8 | 9 | from .attentions import RDCAB 10 | from .misc_gating import ResidualSplitHeadMultiAxisGmlpLayer 11 | 12 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 13 | 14 | 15 | def BottleneckBlock( 16 | features: int, 17 | block_size, 18 | grid_size, 19 | num_groups: int = 1, 20 | block_gmlp_factor: int = 2, 21 | grid_gmlp_factor: int = 2, 22 | input_proj_factor: int = 2, 23 | channels_reduction: int = 4, 24 | dropout_rate: float = 0.0, 25 | use_bias: bool = True, 26 | name: str = "bottleneck_block", 27 | ): 28 | """The bottleneck block consisting of multi-axis gMLP block and RDCAB.""" 29 | 30 | def apply(x): 31 | # input projection 32 | x = Conv1x1(filters=features, use_bias=use_bias, name=f"{name}_input_proj")(x) 33 | shortcut_long = x 34 | 35 | for i in range(num_groups): 36 | x = ResidualSplitHeadMultiAxisGmlpLayer( 37 | grid_size=grid_size, 38 | block_size=block_size, 39 | grid_gmlp_factor=grid_gmlp_factor, 40 | block_gmlp_factor=block_gmlp_factor, 41 | input_proj_factor=input_proj_factor, 42 | use_bias=use_bias, 43 | dropout_rate=dropout_rate, 44 | name=f"{name}_SplitHeadMultiAxisGmlpLayer_{i}", 45 | )(x) 46 | # Channel-mixing part, which provides within-patch communication. 47 | x = RDCAB( 48 | num_channels=features, 49 | reduction=channels_reduction, 50 | use_bias=use_bias, 51 | name=f"{name}_channel_attention_block_1_{i}", 52 | )(x) 53 | 54 | # long skip-connect 55 | x = x + shortcut_long 56 | return x 57 | 58 | return apply 59 | -------------------------------------------------------------------------------- /maxim/blocks/block_gating.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import tensorflow as tf 6 | from tensorflow.keras import backend as K 7 | from tensorflow.keras import layers 8 | 9 | from ..layers import SwapAxes, TFBlockImages, TFUnblockImages 10 | 11 | 12 | def BlockGatingUnit(use_bias: bool = True, name: str = "block_gating_unit"): 13 | """A SpatialGatingUnit as defined in the gMLP paper. 14 | 15 | The 'spatial' dim is defined as the **second last**. 16 | If applied on other dims, you should swapaxes first. 17 | """ 18 | 19 | def apply(x): 20 | u, v = tf.split(x, 2, axis=-1) 21 | v = layers.LayerNormalization( 22 | epsilon=1e-06, name=f"{name}_intermediate_layernorm" 23 | )(v) 24 | n = K.int_shape(x)[-2] # get spatial dim 25 | v = SwapAxes()(v, -1, -2) 26 | v = layers.Dense(n, use_bias=use_bias, name=f"{name}_Dense_0")(v) 27 | v = SwapAxes()(v, -1, -2) 28 | return u * (v + 1.0) 29 | 30 | return apply 31 | 32 | 33 | def BlockGmlpLayer( 34 | block_size, 35 | use_bias: bool = True, 36 | factor: int = 2, 37 | dropout_rate: float = 0.0, 38 | name: str = "block_gmlp", 39 | ): 40 | """Block gMLP layer that performs local mixing of tokens.""" 41 | 42 | def apply(x): 43 | n, h, w, num_channels = ( 44 | K.int_shape(x)[0], 45 | K.int_shape(x)[1], 46 | K.int_shape(x)[2], 47 | K.int_shape(x)[3], 48 | ) 49 | fh, fw = block_size 50 | x, gh, gw = TFBlockImages()(x, patch_size=(fh, fw)) 51 | # MLP2: Local (block) mixing part, provides within-block communication. 52 | y = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm")(x) 53 | y = layers.Dense( 54 | num_channels * factor, 55 | use_bias=use_bias, 56 | name=f"{name}_in_project", 57 | )(y) 58 | y = tf.nn.gelu(y, approximate=True) 59 | y = BlockGatingUnit(use_bias=use_bias, name=f"{name}_BlockGatingUnit")(y) 60 | y = layers.Dense( 61 | num_channels, 62 | use_bias=use_bias, 63 | name=f"{name}_out_project", 64 | )(y) 65 | y = layers.Dropout(dropout_rate)(y) 66 | x = x + y 67 | x = TFUnblockImages()(x, patch_size=(fh, fw), grid_size=(gh, gw)) 68 | return x 69 | 70 | return apply 71 | -------------------------------------------------------------------------------- /maxim/blocks/grid_gating.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import tensorflow as tf 6 | from tensorflow.keras import backend as K 7 | from tensorflow.keras import layers 8 | 9 | from ..layers import SwapAxes, TFBlockImagesByGrid, TFUnblockImages 10 | 11 | 12 | def GridGatingUnit(use_bias: bool = True, name: str = "grid_gating_unit"): 13 | """A SpatialGatingUnit as defined in the gMLP paper. 14 | 15 | The 'spatial' dim is defined as the second last. 16 | If applied on other dims, you should swapaxes first. 17 | """ 18 | 19 | def apply(x): 20 | u, v = tf.split(x, 2, axis=-1) 21 | v = layers.LayerNormalization( 22 | epsilon=1e-06, name=f"{name}_intermediate_layernorm" 23 | )(v) 24 | n = K.int_shape(x)[-3] # get spatial dim 25 | v = SwapAxes()(v, -1, -3) 26 | v = layers.Dense(n, use_bias=use_bias, name=f"{name}_Dense_0")(v) 27 | v = SwapAxes()(v, -1, -3) 28 | return u * (v + 1.0) 29 | 30 | return apply 31 | 32 | 33 | def GridGmlpLayer( 34 | grid_size, 35 | use_bias: bool = True, 36 | factor: int = 2, 37 | dropout_rate: float = 0.0, 38 | name: str = "grid_gmlp", 39 | ): 40 | """Grid gMLP layer that performs global mixing of tokens.""" 41 | 42 | def apply(x): 43 | n, h, w, num_channels = ( 44 | K.int_shape(x)[0], 45 | K.int_shape(x)[1], 46 | K.int_shape(x)[2], 47 | K.int_shape(x)[3], 48 | ) 49 | gh, gw = grid_size 50 | 51 | x, ph, pw = TFBlockImagesByGrid()(x, grid_size=(gh, gw)) 52 | # gMLP1: Global (grid) mixing part, provides global grid communication. 53 | y = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm")(x) 54 | y = layers.Dense( 55 | num_channels * factor, 56 | use_bias=use_bias, 57 | name=f"{name}_in_project", 58 | )(y) 59 | y = tf.nn.gelu(y, approximate=True) 60 | y = GridGatingUnit(use_bias=use_bias, name=f"{name}_GridGatingUnit")(y) 61 | y = layers.Dense( 62 | num_channels, 63 | use_bias=use_bias, 64 | name=f"{name}_out_project", 65 | )(y) 66 | y = layers.Dropout(dropout_rate)(y) 67 | x = x + y 68 | x = TFUnblockImages()(x, grid_size=(gh, gw), patch_size=(ph, pw)) 69 | return x 70 | 71 | return apply 72 | -------------------------------------------------------------------------------- /maxim/tests/conftest.py: -------------------------------------------------------------------------------- 1 | import gc 2 | import random 3 | 4 | import pytest 5 | import tensorflow as tf 6 | from maxim import maxim 7 | from maxim.configs import MAXIM_CONFIGS 8 | from tensorflow import keras 9 | 10 | 11 | @pytest.fixture() 12 | def fix_random(): 13 | tf.random.set_seed(0) 14 | random.seed(0) 15 | 16 | 17 | @pytest.fixture(params=[(16, 12), (12, 16), (16, 16)]) 18 | def window_size(request): 19 | return request.param 20 | 21 | 22 | @pytest.fixture(params=[20, 30, 40]) 23 | def random_image(request, fix_random, window_size): 24 | h, w = window_size 25 | n_windows = request.param 26 | 27 | h_img = h * n_windows 28 | w_img = w * n_windows 29 | 30 | return tf.random.uniform(shape=(5, h_img, w_img, 3), dtype=tf.float32) 31 | 32 | 33 | @pytest.fixture(params=[(10, 13), (14, 15), (20, 20)]) 34 | def random_image_multiple_of_64(request, fix_random, window_size): 35 | h, w = request.param 36 | h_img = h * 64 37 | w_img = w * 64 38 | 39 | return tf.random.uniform(shape=(1, h_img, w_img, 3), dtype=tf.float32) 40 | 41 | 42 | ########################################################################## 43 | ########################## Fixtures for model test ####################### 44 | 45 | 46 | def Model(variant=None, input_resolution=(None, None), **kw) -> keras.Model: 47 | """Factory function to easily create a Model variant like "S". 48 | 49 | Args: 50 | variant: UNet model variants. Options: 'S-1' | 'S-2' | 'S-3' 51 | | 'M-1' | 'M-2' | 'M-3' 52 | input_resolution: Size of the input images. 53 | **kw: Other UNet config dicts. 54 | 55 | Returns: 56 | The MAXIM model. 57 | """ 58 | 59 | if variant is not None: 60 | config = MAXIM_CONFIGS[variant] 61 | for k, v in config.items(): 62 | kw.setdefault(k, v) 63 | 64 | if "variant" in kw: 65 | _ = kw.pop("variant") 66 | if "input_resolution" in kw: 67 | _ = kw.pop("input_resolution") 68 | model_name = kw.pop("name") 69 | 70 | maxim_model = maxim.MAXIM(**kw) 71 | 72 | inputs = keras.Input((*input_resolution, 3)) 73 | outputs = maxim_model(inputs) 74 | final_model = keras.Model(inputs, outputs, name=f"{model_name}_model") 75 | 76 | return final_model 77 | 78 | 79 | # Scope = session means it should only be instantiated once per test session. 80 | @pytest.fixture(scope="session", params=["S-2"]) 81 | def none_model(request): 82 | model = Model(variant=request.param, input_resolution=(None, None)) 83 | yield model 84 | del model 85 | gc.collect() 86 | -------------------------------------------------------------------------------- /maxim/configs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Configs based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | MAXIM_CONFIGS = { 6 | # params: 6.108515000000001 M, GFLOPS: 93.163716608 7 | "S-1": { 8 | "features": 32, 9 | "depth": 3, 10 | "num_stages": 1, 11 | "num_groups": 2, 12 | "num_bottleneck_blocks": 2, 13 | "block_gmlp_factor": 2, 14 | "grid_gmlp_factor": 2, 15 | "input_proj_factor": 2, 16 | "channels_reduction": 4, 17 | "name": "s1", 18 | }, 19 | # params: 13.35383 M, GFLOPS: 206.743273472 20 | "S-2": { 21 | "features": 32, 22 | "depth": 3, 23 | "num_stages": 2, 24 | "num_groups": 2, 25 | "num_bottleneck_blocks": 2, 26 | "block_gmlp_factor": 2, 27 | "grid_gmlp_factor": 2, 28 | "input_proj_factor": 2, 29 | "channels_reduction": 4, 30 | "name": "s2", 31 | }, 32 | # params: 20.599145 M, GFLOPS: 320.32194560000005 33 | "S-3": { 34 | "features": 32, 35 | "depth": 3, 36 | "num_stages": 3, 37 | "num_groups": 2, 38 | "num_bottleneck_blocks": 2, 39 | "block_gmlp_factor": 2, 40 | "grid_gmlp_factor": 2, 41 | "input_proj_factor": 2, 42 | "channels_reduction": 4, 43 | "name": "s3", 44 | }, 45 | # params: 19.361219000000002 M, 308.495712256 GFLOPs 46 | "M-1": { 47 | "features": 64, 48 | "depth": 3, 49 | "num_stages": 1, 50 | "num_groups": 2, 51 | "num_bottleneck_blocks": 2, 52 | "block_gmlp_factor": 2, 53 | "grid_gmlp_factor": 2, 54 | "input_proj_factor": 2, 55 | "channels_reduction": 4, 56 | "name": "m1", 57 | }, 58 | # params: 40.83911 M, 675.25541888 GFLOPs 59 | "M-2": { 60 | "features": 64, 61 | "depth": 3, 62 | "num_stages": 2, 63 | "num_groups": 2, 64 | "num_bottleneck_blocks": 2, 65 | "block_gmlp_factor": 2, 66 | "grid_gmlp_factor": 2, 67 | "input_proj_factor": 2, 68 | "channels_reduction": 4, 69 | "name": "m2", 70 | }, 71 | # params: 62.317001 M, 1042.014666752 GFLOPs 72 | "M-3": { 73 | "features": 64, 74 | "depth": 3, 75 | "num_stages": 3, 76 | "num_groups": 2, 77 | "num_bottleneck_blocks": 2, 78 | "block_gmlp_factor": 2, 79 | "grid_gmlp_factor": 2, 80 | "input_proj_factor": 2, 81 | "channels_reduction": 4, 82 | "name": "m3", 83 | }, 84 | } 85 | -------------------------------------------------------------------------------- /hub_utilities/generate_doc.py: -------------------------------------------------------------------------------- 1 | """Generates model documentation for MAXIM TF models. 2 | 3 | Credits: Willi Gierke 4 | """ 5 | 6 | import os 7 | from string import Template 8 | 9 | import attr 10 | 11 | template = Template( 12 | """# Module $HANDLE 13 | 14 | MAXIM model pre-trained on the $DATASET_DESCRIPTION suitable for image $TASK. 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | ## Overview 24 | 25 | This model is based on the MAXIM backbone [1] pre-trained on the $DATASET_DESCRIPTION. You can use this 26 | model for image $TASK. Please refer to the Colab Notebook linked on this page for more details. 27 | 28 | MAXIM introduces a common backbone for different image processing tasks like 29 | denoising, deblurring, dehazing, deraining, and enhancement. You can find the complete 30 | collection of MAXIM models on TF-Hub on [this page](https://tfhub.dev/sayakpaul/collections/maxim/1). 31 | 32 | ## Notes 33 | 34 | * The original model weights are provided in [2]. There were ported to Keras models 35 | (`tf.keras.Model`) and then serialized as TensorFlow SavedModels. The porting 36 | steps are available in [3]. 37 | * The format of the model handle is: `'maxim_{variant}_{task}_{dataset}'`. 38 | * The model can be unrolled into a standard Keras model and you can inspect its topology. 39 | To do so, first download the model from TF-Hub and then load it using `tf.keras.models.load_model` 40 | providing the path to the downloaded model folder. 41 | 42 | ## References 43 | 44 | [1] [MAXIM: Multi-Axis MLP for Image Processing Tu et al.](https://arxiv.org/abs/2201.02973) 45 | 46 | [2] [MAXIM GitHub](https://github.com/google-research/maxim) 47 | 48 | [3] [MAXIM TF GitHub](https://github.com/sayakpaul/maxim-tf) 49 | 50 | ## Acknowledgements 51 | 52 | * [Gustavo Martins](https://twitter.com/gusthema?lang=en) 53 | * [ML-GDE program](https://developers.google.com/programs/experts/) 54 | 55 | """ 56 | ) 57 | 58 | 59 | @attr.s 60 | class Config: 61 | variant = attr.ib(type=str) 62 | dataset = attr.ib(type=str) 63 | task = attr.ib(type=str) 64 | task_metadata = attr.ib(type=str) 65 | 66 | def gcs_folder_name(self): 67 | return f"{self.variant}_{self.task}_{self.dataset}" 68 | 69 | def handle(self): 70 | return f"sayakpaul/maxim_{self.gcs_folder_name().lower()}/1" 71 | 72 | def rel_doc_file_path(self): 73 | """Relative to the tfhub.dev directory.""" 74 | return f"assets/docs/{self.handle()}.md" 75 | 76 | 77 | for c in [ 78 | Config("S-2", "sots-indoor", "dehazing", "dehazing"), 79 | Config("S-2", "sots-outdoor", "dehazing", "dehazing"), 80 | Config("S-2", "rain13k", "deraining", "deraining"), 81 | Config("S-2", "raindrop", "deraining", "deraining"), 82 | Config("S-2", "fivek", "enhancement", "enhancement"), 83 | Config("S-2", "lol", "enhancement", "enhancement"), 84 | Config("S-3", "gopro", "deblurring", "deblurring"), 85 | Config("S-3", "realblur_j", "deblurring", "deblurring"), 86 | Config("S-3", "realblur_r", "deblurring", "deblurring"), 87 | Config("S-3", "reds", "deblurring", "deblurring"), 88 | Config("S-3", "sidd", "denoising", "denoising"), 89 | ]: 90 | save_path = os.path.join( 91 | "/Users/sayakpaul/Downloads/", "tfhub.dev", c.rel_doc_file_path() 92 | ) 93 | model_folder = save_path.split("/")[-2] 94 | model_abs_path = "/".join(save_path.split("/")[:-1]) 95 | 96 | if not os.path.exists(model_abs_path): 97 | os.makedirs(model_abs_path, exist_ok=True) 98 | 99 | with open(save_path, "w") as f: 100 | f.write( 101 | template.substitute( 102 | HANDLE=c.handle(), 103 | DATASET_DESCRIPTION=c.dataset, 104 | TASK=c.task, 105 | TASK_METADATA=c.task_metadata, 106 | ARCHIVE_NAME=c.gcs_folder_name(), 107 | ) 108 | ) 109 | -------------------------------------------------------------------------------- /maxim/blocks/unet.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | import tensorflow as tf 8 | from tensorflow.keras import layers 9 | 10 | from .attentions import RCAB 11 | from .misc_gating import CrossGatingBlock, ResidualSplitHeadMultiAxisGmlpLayer 12 | 13 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 14 | Conv3x3 = functools.partial(layers.Conv2D, kernel_size=(3, 3), padding="same") 15 | ConvT_up = functools.partial( 16 | layers.Conv2DTranspose, kernel_size=(2, 2), strides=(2, 2), padding="same" 17 | ) 18 | Conv_down = functools.partial( 19 | layers.Conv2D, kernel_size=(4, 4), strides=(2, 2), padding="same" 20 | ) 21 | 22 | 23 | def UNetEncoderBlock( 24 | num_channels: int, 25 | block_size, 26 | grid_size, 27 | num_groups: int = 1, 28 | lrelu_slope: float = 0.2, 29 | block_gmlp_factor: int = 2, 30 | grid_gmlp_factor: int = 2, 31 | input_proj_factor: int = 2, 32 | channels_reduction: int = 4, 33 | dropout_rate: float = 0.0, 34 | downsample: bool = True, 35 | use_global_mlp: bool = True, 36 | use_bias: bool = True, 37 | use_cross_gating: bool = False, 38 | name: str = "unet_encoder", 39 | ): 40 | """Encoder block in MAXIM.""" 41 | 42 | def apply(x, skip=None, enc=None, dec=None): 43 | if skip is not None: 44 | x = tf.concat([x, skip], axis=-1) 45 | 46 | # convolution-in 47 | x = Conv1x1(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_0")(x) 48 | shortcut_long = x 49 | 50 | for i in range(num_groups): 51 | if use_global_mlp: 52 | x = ResidualSplitHeadMultiAxisGmlpLayer( 53 | grid_size=grid_size, 54 | block_size=block_size, 55 | grid_gmlp_factor=grid_gmlp_factor, 56 | block_gmlp_factor=block_gmlp_factor, 57 | input_proj_factor=input_proj_factor, 58 | use_bias=use_bias, 59 | dropout_rate=dropout_rate, 60 | name=f"{name}_SplitHeadMultiAxisGmlpLayer_{i}", 61 | )(x) 62 | x = RCAB( 63 | num_channels=num_channels, 64 | reduction=channels_reduction, 65 | lrelu_slope=lrelu_slope, 66 | use_bias=use_bias, 67 | name=f"{name}_channel_attention_block_1{i}", 68 | )(x) 69 | 70 | x = x + shortcut_long 71 | 72 | if enc is not None and dec is not None: 73 | assert use_cross_gating 74 | x, _ = CrossGatingBlock( 75 | features=num_channels, 76 | block_size=block_size, 77 | grid_size=grid_size, 78 | dropout_rate=dropout_rate, 79 | input_proj_factor=input_proj_factor, 80 | upsample_y=False, 81 | use_bias=use_bias, 82 | name=f"{name}_cross_gating_block", 83 | )(x, enc + dec) 84 | 85 | if downsample: 86 | x_down = Conv_down( 87 | filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_1" 88 | )(x) 89 | return x_down, x 90 | else: 91 | return x 92 | 93 | return apply 94 | 95 | 96 | def UNetDecoderBlock( 97 | num_channels: int, 98 | block_size, 99 | grid_size, 100 | num_groups: int = 1, 101 | lrelu_slope: float = 0.2, 102 | block_gmlp_factor: int = 2, 103 | grid_gmlp_factor: int = 2, 104 | input_proj_factor: int = 2, 105 | channels_reduction: int = 4, 106 | dropout_rate: float = 0.0, 107 | downsample: bool = True, 108 | use_global_mlp: bool = True, 109 | use_bias: bool = True, 110 | name: str = "unet_decoder", 111 | ): 112 | 113 | """Decoder block in MAXIM.""" 114 | 115 | def apply(x, bridge=None): 116 | x = ConvT_up( 117 | filters=num_channels, use_bias=use_bias, name=f"{name}_ConvTranspose_0" 118 | )(x) 119 | x = UNetEncoderBlock( 120 | num_channels=num_channels, 121 | num_groups=num_groups, 122 | lrelu_slope=lrelu_slope, 123 | block_size=block_size, 124 | grid_size=grid_size, 125 | block_gmlp_factor=block_gmlp_factor, 126 | grid_gmlp_factor=grid_gmlp_factor, 127 | channels_reduction=channels_reduction, 128 | use_global_mlp=use_global_mlp, 129 | dropout_rate=dropout_rate, 130 | downsample=False, 131 | use_bias=use_bias, 132 | name=f"{name}_UNetEncoderBlock_0", 133 | )(x, skip=bridge) 134 | 135 | return x 136 | 137 | return apply 138 | -------------------------------------------------------------------------------- /maxim/blocks/attentions.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | import tensorflow as tf 8 | from tensorflow.keras import layers 9 | 10 | from .others import MlpBlock 11 | 12 | Conv3x3 = functools.partial(layers.Conv2D, kernel_size=(3, 3), padding="same") 13 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 14 | 15 | 16 | def CALayer( 17 | num_channels: int, 18 | reduction: int = 4, 19 | use_bias: bool = True, 20 | name: str = "channel_attention", 21 | ): 22 | """Squeeze-and-excitation block for channel attention. 23 | 24 | ref: https://arxiv.org/abs/1709.01507 25 | """ 26 | 27 | def apply(x): 28 | # 2D global average pooling 29 | y = layers.GlobalAvgPool2D(keepdims=True)(x) 30 | # Squeeze (in Squeeze-Excitation) 31 | y = Conv1x1( 32 | filters=num_channels // reduction, use_bias=use_bias, name=f"{name}_Conv_0" 33 | )(y) 34 | y = tf.nn.relu(y) 35 | # Excitation (in Squeeze-Excitation) 36 | y = Conv1x1(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_1")(y) 37 | y = tf.nn.sigmoid(y) 38 | return x * y 39 | 40 | return apply 41 | 42 | 43 | def RCAB( 44 | num_channels: int, 45 | reduction: int = 4, 46 | lrelu_slope: float = 0.2, 47 | use_bias: bool = True, 48 | name: str = "residual_ca", 49 | ): 50 | """Residual channel attention block. Contains LN,Conv,lRelu,Conv,SELayer.""" 51 | 52 | def apply(x): 53 | shortcut = x 54 | x = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm")(x) 55 | x = Conv3x3(filters=num_channels, use_bias=use_bias, name=f"{name}_conv1")(x) 56 | x = tf.nn.leaky_relu(x, alpha=lrelu_slope) 57 | x = Conv3x3(filters=num_channels, use_bias=use_bias, name=f"{name}_conv2")(x) 58 | x = CALayer( 59 | num_channels=num_channels, 60 | reduction=reduction, 61 | use_bias=use_bias, 62 | name=f"{name}_channel_attention", 63 | )(x) 64 | return x + shortcut 65 | 66 | return apply 67 | 68 | 69 | def RDCAB( 70 | num_channels: int, 71 | reduction: int = 16, 72 | use_bias: bool = True, 73 | dropout_rate: float = 0.0, 74 | name: str = "rdcab", 75 | ): 76 | """Residual dense channel attention block. Used in Bottlenecks.""" 77 | 78 | def apply(x): 79 | y = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm")(x) 80 | y = MlpBlock( 81 | mlp_dim=num_channels, 82 | dropout_rate=dropout_rate, 83 | use_bias=use_bias, 84 | name=f"{name}_channel_mixing", 85 | )(y) 86 | y = CALayer( 87 | num_channels=num_channels, 88 | reduction=reduction, 89 | use_bias=use_bias, 90 | name=f"{name}_channel_attention", 91 | )(y) 92 | x = x + y 93 | return x 94 | 95 | return apply 96 | 97 | 98 | def SAM( 99 | num_channels: int, 100 | output_channels: int = 3, 101 | use_bias: bool = True, 102 | name: str = "sam", 103 | ): 104 | 105 | """Supervised attention module for multi-stage training. 106 | 107 | Introduced by MPRNet [CVPR2021]: https://github.com/swz30/MPRNet 108 | """ 109 | 110 | def apply(x, x_image): 111 | """Apply the SAM module to the input and num_channels. 112 | Args: 113 | x: the output num_channels from UNet decoder with shape (h, w, c) 114 | x_image: the input image with shape (h, w, 3) 115 | Returns: 116 | A tuple of tensors (x1, image) where (x1) is the sam num_channels used for the 117 | next stage, and (image) is the output restored image at current stage. 118 | """ 119 | # Get num_channels 120 | x1 = Conv3x3(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_0")(x) 121 | 122 | # Output restored image X_s 123 | if output_channels == 3: 124 | image = ( 125 | Conv3x3( 126 | filters=output_channels, use_bias=use_bias, name=f"{name}_Conv_1" 127 | )(x) 128 | + x_image 129 | ) 130 | else: 131 | image = Conv3x3( 132 | filters=output_channels, use_bias=use_bias, name=f"{name}_Conv_1" 133 | )(x) 134 | 135 | # Get attention maps for num_channels 136 | x2 = tf.nn.sigmoid( 137 | Conv3x3(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_2")( 138 | image 139 | ) 140 | ) 141 | 142 | # Get attended feature maps 143 | x1 = x1 * x2 144 | 145 | # Residual connection 146 | x1 = x1 + x 147 | return x1, image 148 | 149 | return apply 150 | -------------------------------------------------------------------------------- /maxim/layers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Layers based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | and reworked to cope with variable image dimensions 4 | """ 5 | 6 | import tensorflow as tf 7 | from tensorflow.experimental import numpy as tnp 8 | from tensorflow.keras import layers 9 | 10 | 11 | @tf.keras.utils.register_keras_serializable("maxim") 12 | class TFBlockImages(layers.Layer): 13 | def __init__(self, **kwargs): 14 | super().__init__(**kwargs) 15 | 16 | def call(self, image, patch_size): 17 | bs, h, w, num_channels = ( 18 | tf.shape(image)[0], 19 | tf.shape(image)[1], 20 | tf.shape(image)[2], 21 | tf.shape(image)[3], 22 | ) 23 | ph, pw = patch_size 24 | gh = h // ph 25 | gw = w // pw 26 | pad = [[0, 0], [0, 0]] 27 | patches = tf.space_to_batch_nd(image, [ph, pw], pad) 28 | patches = tf.split(patches, ph * pw, axis=0) 29 | patches = tf.stack(patches, 3) # (bs, h/p, h/p, p*p, 3) 30 | patches_dim = tf.shape(patches) 31 | patches = tf.reshape( 32 | patches, [patches_dim[0], patches_dim[1], patches_dim[2], -1] 33 | ) 34 | patches = tf.reshape( 35 | patches, 36 | (patches_dim[0], patches_dim[1] * patches_dim[2], ph * pw, num_channels), 37 | ) 38 | return [patches, gh, gw] 39 | 40 | def get_config(self): 41 | return super().get_config() 42 | 43 | 44 | @tf.keras.utils.register_keras_serializable("maxim") 45 | class TFBlockImagesByGrid(layers.Layer): 46 | def __init__(self, **kwargs): 47 | super().__init__(**kwargs) 48 | 49 | def call(self, image, grid_size): 50 | bs, h, w, num_channels = ( 51 | tf.shape(image)[0], 52 | tf.shape(image)[1], 53 | tf.shape(image)[2], 54 | tf.shape(image)[3], 55 | ) 56 | gh, gw = grid_size 57 | ph = h // gh 58 | pw = w // gw 59 | pad = [[0, 0], [0, 0]] 60 | 61 | def block_single_image(img): 62 | pat = tf.expand_dims(img, 0) # batch = 1 63 | pat = tf.space_to_batch_nd(pat, [ph, pw], pad) # p*p*bs, g, g, c 64 | pat = tf.expand_dims(pat, 3) # pxpxbs, g, g, 1, c 65 | pat = tf.transpose(pat, perm=[3, 1, 2, 0, 4]) # 1, g, g, pxp, c 66 | pat = tf.reshape(pat, [gh, gw, ph * pw, num_channels]) 67 | return pat 68 | 69 | patches = image 70 | patches = tf.map_fn(fn=lambda x: block_single_image(x), elems=patches) 71 | patches_dim = tf.shape(patches) 72 | patches = tf.reshape( 73 | patches, [patches_dim[0], patches_dim[1], patches_dim[2], -1] 74 | ) 75 | patches = tf.reshape( 76 | patches, 77 | (patches_dim[0], patches_dim[1] * patches_dim[2], ph * pw, num_channels), 78 | ) 79 | return [patches, ph, pw] 80 | 81 | def get_config(self): 82 | return super().get_config() 83 | 84 | 85 | @tf.keras.utils.register_keras_serializable("maxim") 86 | class TFUnblockImages(layers.Layer): 87 | def __init__(self, **kwargs): 88 | super().__init__(**kwargs) 89 | 90 | def call(self, x, patch_size, grid_size): 91 | bs, grid_sqrt, patch_sqrt, num_channels = ( 92 | tf.shape(x)[0], 93 | tf.shape(x)[1], 94 | tf.shape(x)[2], 95 | tf.shape(x)[3], 96 | ) 97 | ph, pw = patch_size 98 | gh, gw = grid_size 99 | 100 | pad = [[0, 0], [0, 0]] 101 | 102 | y = tf.reshape(x, (bs, gh, gw, -1, num_channels)) # (bs, gh, gw, ph*pw, 3) 103 | y = tf.expand_dims(y, 0) 104 | y = tf.transpose(y, perm=[4, 1, 2, 3, 0, 5]) 105 | y = tf.reshape(y, [bs * ph * pw, gh, gw, num_channels]) 106 | y = tf.batch_to_space(y, [ph, pw], pad) 107 | 108 | return y 109 | 110 | def get_config(self): 111 | return super().get_config() 112 | 113 | 114 | @tf.keras.utils.register_keras_serializable("maxim") 115 | class SwapAxes(layers.Layer): 116 | def __init__(self, **kwargs): 117 | super().__init__(**kwargs) 118 | 119 | def call(self, x, axis_one, axis_two): 120 | return tnp.swapaxes(x, axis_one, axis_two) 121 | 122 | def get_config(self): 123 | config = super().get_config().copy() 124 | return config 125 | 126 | 127 | @tf.keras.utils.register_keras_serializable("maxim") 128 | class Resizing(tf.keras.layers.Layer): 129 | def __init__(self, ratio: float, method="bilinear", antialias=True, **kwargs): 130 | super().__init__(**kwargs) 131 | self.ratio = ratio 132 | self.method = method 133 | self.antialias = antialias 134 | 135 | def call(self, img): 136 | shape = tf.shape(img) 137 | 138 | new_sh = tf.cast(shape[1:3], tf.float32) // self.ratio 139 | 140 | x = tf.image.resize( 141 | img, 142 | size=tf.cast(new_sh, tf.int32), 143 | method=self.method, 144 | antialias=self.antialias, 145 | ) 146 | return x 147 | 148 | def get_config(self): 149 | config = super().get_config().copy() 150 | config.update( 151 | { 152 | "ratio": self.ratio, 153 | "antialias": self.antialias, 154 | "method": self.method, 155 | } 156 | ) 157 | return config 158 | -------------------------------------------------------------------------------- /maxim/tests/test_block_operations.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import einops 4 | import numpy as np 5 | import tensorflow as tf 6 | from maxim.layers import TFBlockImages, TFBlockImagesByGrid, TFUnblockImages 7 | from tensorflow.keras import backend as K 8 | from tensorflow.keras import layers 9 | 10 | LOW_THRESHOLD = 1e-7 11 | 12 | 13 | @tf.keras.utils.register_keras_serializable("maxim") 14 | class BlockImages(layers.Layer): 15 | def __init__(self, **kwargs): 16 | super().__init__(**kwargs) 17 | 18 | def call(self, x, patch_size): 19 | bs, h, w, num_channels = ( 20 | K.int_shape(x)[0], 21 | K.int_shape(x)[1], 22 | K.int_shape(x)[2], 23 | K.int_shape(x)[3], 24 | ) 25 | 26 | grid_height, grid_width = h // patch_size[0], w // patch_size[1] 27 | 28 | x = einops.rearrange( 29 | x, 30 | "n (gh fh) (gw fw) c -> n (gh gw) (fh fw) c", 31 | gh=grid_height, 32 | gw=grid_width, 33 | fh=patch_size[0], 34 | fw=patch_size[1], 35 | ) 36 | 37 | return x 38 | 39 | def get_config(self): 40 | config = super().get_config().copy() 41 | return config 42 | 43 | 44 | @tf.keras.utils.register_keras_serializable("maxim") 45 | class UnblockImages(layers.Layer): 46 | def __init__(self, **kwargs): 47 | super().__init__(**kwargs) 48 | 49 | def call(self, x, grid_size, patch_size): 50 | x = einops.rearrange( 51 | x, 52 | "n (gh gw) (fh fw) c -> n (gh fh) (gw fw) c", 53 | gh=grid_size[0], 54 | gw=grid_size[1], 55 | fh=patch_size[0], 56 | fw=patch_size[1], 57 | ) 58 | 59 | return x 60 | 61 | def get_config(self): 62 | config = super().get_config().copy() 63 | return config 64 | 65 | 66 | def test_patch_block_equivalence(random_image, window_size): 67 | patch_size = window_size 68 | patched_image_original = BlockImages()(random_image, patch_size=patch_size) 69 | patched_image_tf, _, _ = TFBlockImages()(random_image, patch_size=patch_size) 70 | difference = np.sum( 71 | (patched_image_original.numpy() - patched_image_tf.numpy()) ** 2 72 | ) 73 | assert difference < LOW_THRESHOLD 74 | 75 | 76 | def test_grid_block_equivalence(random_image, window_size): 77 | grid_size = window_size 78 | gh, gw = grid_size 79 | height, width = random_image.shape[1], random_image.shape[2] 80 | patch_size = (height // gh, width // gw) 81 | patched_image_original = BlockImages()(random_image, patch_size=patch_size) 82 | patched_image_tf, _, _ = TFBlockImagesByGrid()(random_image, grid_size=grid_size) 83 | difference = np.sum( 84 | (patched_image_original.numpy() - patched_image_tf.numpy()) ** 2 85 | ) 86 | assert difference < LOW_THRESHOLD 87 | 88 | 89 | def test_reconstruction_by_grid(random_image, window_size): 90 | grid_size = window_size 91 | height, width = random_image.shape[1], random_image.shape[2] 92 | p_h, p_w = height // grid_size[0], width // grid_size[1] 93 | 94 | # Block and Unblock with einops layers 95 | patched_image_original = BlockImages()(random_image, patch_size=(p_h, p_w)) 96 | reconstructed_original = UnblockImages()( 97 | patched_image_original, 98 | grid_size=grid_size, 99 | patch_size=(p_h, p_w), 100 | ) 101 | 102 | # Block and Unblock with TF layers 103 | patched_image_tf, ph, pw = TFBlockImagesByGrid()(random_image, grid_size=grid_size) 104 | reconstructed_image_tf = TFUnblockImages()( 105 | patched_image_tf, grid_size=grid_size, patch_size=(ph, pw) 106 | ) 107 | 108 | # Compare implementation diff and reconstruction diff 109 | difference_between_implementations = np.sum( 110 | (reconstructed_original.numpy() - reconstructed_image_tf.numpy()) ** 2 111 | ) 112 | assert difference_between_implementations < LOW_THRESHOLD 113 | difference_between_reconstruction = np.sum( 114 | (random_image.numpy() - reconstructed_image_tf.numpy()) ** 2 115 | ) 116 | assert difference_between_reconstruction < LOW_THRESHOLD 117 | 118 | 119 | def test_reconstruction(random_image, window_size): 120 | patch_size = window_size 121 | height, width = random_image.shape[1], random_image.shape[2] 122 | 123 | # Block and Unblock with einops layers 124 | patched_image_original = BlockImages()(random_image, patch_size=patch_size) 125 | reconstructed_original = UnblockImages()( 126 | patched_image_original, 127 | grid_size=(height // patch_size[0], width // patch_size[1]), 128 | patch_size=patch_size, 129 | ) 130 | 131 | # Block and Unblock with TF layers 132 | patched_image_tf, gh, gw = TFBlockImages()(random_image, patch_size=patch_size) 133 | reconstructed_image_tf = TFUnblockImages()( 134 | patched_image_tf, patch_size=patch_size, grid_size=(gh, gw) 135 | ) 136 | 137 | # Compare implementation diff and reconstruction diff 138 | difference_between_implementations = np.sum( 139 | (reconstructed_original.numpy() - reconstructed_image_tf.numpy()) ** 2 140 | ) 141 | assert difference_between_implementations < LOW_THRESHOLD 142 | difference_between_reconstruction = np.sum( 143 | (random_image.numpy() - reconstructed_image_tf.numpy()) ** 2 144 | ) 145 | assert difference_between_reconstruction < LOW_THRESHOLD 146 | -------------------------------------------------------------------------------- /maxim/blocks/misc_gating.py: -------------------------------------------------------------------------------- 1 | """ 2 | Blocks based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | import tensorflow as tf 8 | from tensorflow.keras import backend as K 9 | from tensorflow.keras import layers 10 | 11 | from ..layers import SwapAxes, TFBlockImages, TFBlockImagesByGrid, TFUnblockImages 12 | from .block_gating import BlockGmlpLayer 13 | from .grid_gating import GridGmlpLayer 14 | 15 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 16 | Conv3x3 = functools.partial(layers.Conv2D, kernel_size=(3, 3), padding="same") 17 | ConvT_up = functools.partial( 18 | layers.Conv2DTranspose, kernel_size=(2, 2), strides=(2, 2), padding="same" 19 | ) 20 | Conv_down = functools.partial( 21 | layers.Conv2D, kernel_size=(4, 4), strides=(2, 2), padding="same" 22 | ) 23 | 24 | 25 | def ResidualSplitHeadMultiAxisGmlpLayer( 26 | block_size, 27 | grid_size, 28 | block_gmlp_factor: int = 2, 29 | grid_gmlp_factor: int = 2, 30 | input_proj_factor: int = 2, 31 | use_bias: bool = True, 32 | dropout_rate: float = 0.0, 33 | name: str = "residual_split_head_maxim", 34 | ): 35 | """The multi-axis gated MLP block.""" 36 | 37 | def apply(x): 38 | shortcut = x 39 | n, h, w, num_channels = ( 40 | K.int_shape(x)[0], 41 | K.int_shape(x)[1], 42 | K.int_shape(x)[2], 43 | K.int_shape(x)[3], 44 | ) 45 | x = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm_in")(x) 46 | 47 | x = layers.Dense( 48 | int(num_channels) * input_proj_factor, 49 | use_bias=use_bias, 50 | name=f"{name}_in_project", 51 | )(x) 52 | x = tf.nn.gelu(x, approximate=True) 53 | 54 | u, v = tf.split(x, 2, axis=-1) 55 | 56 | # GridGMLPLayer 57 | u = GridGmlpLayer( 58 | grid_size=grid_size, 59 | factor=grid_gmlp_factor, 60 | use_bias=use_bias, 61 | dropout_rate=dropout_rate, 62 | name=f"{name}_GridGmlpLayer", 63 | )(u) 64 | 65 | # BlockGMLPLayer 66 | v = BlockGmlpLayer( 67 | block_size=block_size, 68 | factor=block_gmlp_factor, 69 | use_bias=use_bias, 70 | dropout_rate=dropout_rate, 71 | name=f"{name}_BlockGmlpLayer", 72 | )(v) 73 | 74 | x = tf.concat([u, v], axis=-1) 75 | 76 | x = layers.Dense( 77 | num_channels, 78 | use_bias=use_bias, 79 | name=f"{name}_out_project", 80 | )(x) 81 | x = layers.Dropout(dropout_rate)(x) 82 | x = x + shortcut 83 | return x 84 | 85 | return apply 86 | 87 | 88 | def GetSpatialGatingWeights( 89 | features: int, 90 | block_size, 91 | grid_size, 92 | input_proj_factor: int = 2, 93 | dropout_rate: float = 0.0, 94 | use_bias: bool = True, 95 | name: str = "spatial_gating", 96 | ): 97 | 98 | """Get gating weights for cross-gating MLP block.""" 99 | 100 | def apply(x): 101 | n, h, w, num_channels = ( 102 | K.int_shape(x)[0], 103 | K.int_shape(x)[1], 104 | K.int_shape(x)[2], 105 | K.int_shape(x)[3], 106 | ) 107 | 108 | # input projection 109 | x = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm_in")(x) 110 | x = layers.Dense( 111 | num_channels * input_proj_factor, 112 | use_bias=use_bias, 113 | name=f"{name}_in_project", 114 | )(x) 115 | x = tf.nn.gelu(x, approximate=True) 116 | u, v = tf.split(x, 2, axis=-1) 117 | 118 | # Get grid MLP weights 119 | gh, gw = grid_size 120 | u, phu, pwu = TFBlockImagesByGrid()(u, grid_size=(gh, gw)) 121 | dim_u = gh * gw 122 | u = SwapAxes()(u, -1, -3) 123 | u = layers.Dense(dim_u, use_bias=use_bias, name=f"{name}_Dense_0")(u) 124 | u = SwapAxes()(u, -1, -3) 125 | u = TFUnblockImages()(u, grid_size=(gh, gw), patch_size=(phu, pwu)) 126 | 127 | # Get Block MLP weights 128 | fh, fw = block_size 129 | v, gh, gw = TFBlockImages()(v, patch_size=(fh, fw)) 130 | dim_v = fh * fw 131 | v = SwapAxes()(v, -1, -2) 132 | v = layers.Dense(dim_v, use_bias=use_bias, name=f"{name}_Dense_1")(v) 133 | v = SwapAxes()(v, -1, -2) 134 | v = TFUnblockImages()(v, patch_size=(fh, fw), grid_size=(gh, gw)) 135 | 136 | x = tf.concat([u, v], axis=-1) 137 | x = layers.Dense(num_channels, use_bias=use_bias, name=f"{name}_out_project")(x) 138 | x = layers.Dropout(dropout_rate)(x) 139 | return x 140 | 141 | return apply 142 | 143 | 144 | def CrossGatingBlock( 145 | features: int, 146 | block_size, 147 | grid_size, 148 | dropout_rate: float = 0.0, 149 | input_proj_factor: int = 2, 150 | upsample_y: bool = True, 151 | use_bias: bool = True, 152 | name: str = "cross_gating", 153 | ): 154 | 155 | """Cross-gating MLP block.""" 156 | 157 | def apply(x, y): 158 | # Upscale Y signal, y is the gating signal. 159 | if upsample_y: 160 | y = ConvT_up( 161 | filters=features, use_bias=use_bias, name=f"{name}_ConvTranspose_0" 162 | )(y) 163 | 164 | x = Conv1x1(filters=features, use_bias=use_bias, name=f"{name}_Conv_0")(x) 165 | n, h, w, num_channels = ( 166 | K.int_shape(x)[0], 167 | K.int_shape(x)[1], 168 | K.int_shape(x)[2], 169 | K.int_shape(x)[3], 170 | ) 171 | 172 | y = Conv1x1(filters=num_channels, use_bias=use_bias, name=f"{name}_Conv_1")(y) 173 | 174 | shortcut_x = x 175 | shortcut_y = y 176 | 177 | # Get gating weights from X 178 | x = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm_x")(x) 179 | x = layers.Dense(num_channels, use_bias=use_bias, name=f"{name}_in_project_x")( 180 | x 181 | ) 182 | x = tf.nn.gelu(x, approximate=True) 183 | gx = GetSpatialGatingWeights( 184 | features=num_channels, 185 | block_size=block_size, 186 | grid_size=grid_size, 187 | dropout_rate=dropout_rate, 188 | use_bias=use_bias, 189 | name=f"{name}_SplitHeadMultiAxisGating_x", 190 | )(x) 191 | 192 | # Get gating weights from Y 193 | y = layers.LayerNormalization(epsilon=1e-06, name=f"{name}_LayerNorm_y")(y) 194 | y = layers.Dense(num_channels, use_bias=use_bias, name=f"{name}_in_project_y")( 195 | y 196 | ) 197 | y = tf.nn.gelu(y, approximate=True) 198 | gy = GetSpatialGatingWeights( 199 | features=num_channels, 200 | block_size=block_size, 201 | grid_size=grid_size, 202 | dropout_rate=dropout_rate, 203 | use_bias=use_bias, 204 | name=f"{name}_SplitHeadMultiAxisGating_y", 205 | )(y) 206 | 207 | # Apply cross gating: X = X * GY, Y = Y * GX 208 | y = y * gx 209 | y = layers.Dense(num_channels, use_bias=use_bias, name=f"{name}_out_project_y")( 210 | y 211 | ) 212 | y = layers.Dropout(dropout_rate)(y) 213 | y = y + shortcut_y 214 | 215 | x = x * gy # gating x using y 216 | x = layers.Dense(num_channels, use_bias=use_bias, name=f"{name}_out_project_x")( 217 | x 218 | ) 219 | x = layers.Dropout(dropout_rate)(x) 220 | x = x + y + shortcut_x # get all aggregated signals 221 | return x, y 222 | 223 | return apply 224 | -------------------------------------------------------------------------------- /convert_to_tf.py: -------------------------------------------------------------------------------- 1 | """ 2 | Script to port the pre-trained JAX params of MAXIM to TF. 3 | 4 | Usage: 5 | python convert_to_tf.py 6 | 7 | The above will convert a MAXIM-3S model trained on the denoising task with the 8 | SIDD dataset. You can find the tasks and checkpoints supported by MAXIM here: 9 | https://github.com/google-research/maxim#results-and-pre-trained-models. 10 | 11 | So, to convert the pre-trained JAX params (for deblurring on GoPro dataset, say) to TF, 12 | you can run the following: 13 | 14 | python convert_to_tf.py \ 15 | --task Deblurring \ 16 | --ckpt_path gs://gresearch/maxim/ckpt/Deblurring/GoPro/checkpoint.npz 17 | 18 | """ 19 | 20 | import argparse 21 | import collections 22 | import io 23 | import re 24 | from typing import Tuple 25 | 26 | import numpy as np 27 | import pandas as pd 28 | import tensorflow as tf 29 | from huggingface_hub import push_to_hub_keras 30 | 31 | from create_maxim_model import Model 32 | from maxim.configs import MAXIM_CONFIGS 33 | 34 | _MODEL_VARIANT_DICT = { 35 | "Denoising": "S-3", 36 | "Deblurring": "S-3", 37 | "Deraining": "S-2", 38 | "Dehazing": "S-2", 39 | "Enhancement": "S-2", 40 | } 41 | 42 | 43 | # `recover_tree()` and `get_params()` come from here: 44 | # https://github.com/google-research/maxim/blob/main/maxim/run_eval.py 45 | def recover_tree(keys, values): 46 | """Recovers a tree as a nested dict from flat names and values. 47 | This function is useful to analyze checkpoints that are saved by our programs 48 | without need to access the exact source code of the experiment. In particular, 49 | it can be used to extract an reuse various subtrees of the scheckpoint, e.g. 50 | subtree of parameters. 51 | Args: 52 | keys: a list of keys, where '/' is used as separator between nodes. 53 | values: a list of leaf values. 54 | Returns: 55 | A nested tree-like dict. 56 | """ 57 | tree = {} 58 | sub_trees = collections.defaultdict(list) 59 | for k, v in zip(keys, values): 60 | if "/" not in k: 61 | tree[k] = v 62 | else: 63 | k_left, k_right = k.split("/", 1) 64 | sub_trees[k_left].append((k_right, v)) 65 | for k, kv_pairs in sub_trees.items(): 66 | k_subtree, v_subtree = zip(*kv_pairs) 67 | tree[k] = recover_tree(k_subtree, v_subtree) 68 | return tree 69 | 70 | 71 | def get_params(ckpt_path): 72 | """Get params checkpoint.""" 73 | with tf.io.gfile.GFile(ckpt_path, "rb") as f: 74 | data = f.read() 75 | values = np.load(io.BytesIO(data)) 76 | params = recover_tree(*zip(*values.items())) 77 | params = params["opt"]["target"] 78 | return params 79 | 80 | 81 | # From https://stackoverflow.com/questions/5491913/sorting-list-in-python 82 | def sort_nicely(l): 83 | """Sort the given iterable in the way that humans expect.""" 84 | convert = lambda text: int(text) if text.isdigit() else text 85 | alphanum_key = lambda key: [convert(c) for c in re.split("([0-9]+)", key)] 86 | return sorted(l, key=alphanum_key) 87 | 88 | 89 | def modify_upsample(jax_params): 90 | modified_jax_params = collections.OrderedDict() 91 | 92 | jax_keys = list(jax_params.keys()) 93 | keys_upsampling = [] 94 | for k in range(len(jax_keys)): 95 | if "UpSample" in jax_keys[k]: 96 | keys_upsampling.append(jax_keys[k]) 97 | sorted_keys_upsampling = sort_nicely(keys_upsampling) 98 | 99 | i = 1 100 | for k in sorted_keys_upsampling: 101 | k_t = k.split("_")[0] + "_" + str(i) 102 | i += 1 103 | for j in jax_params[k]: 104 | for l in jax_params[k][j]: 105 | modified_param_name = f"{k_t}_{j}/{l}:0" 106 | params = jax_params[k][j][l] 107 | modified_jax_params.update({modified_param_name: params}) 108 | 109 | return modified_jax_params 110 | 111 | 112 | def modify_jax_params(jax_params): 113 | modified_jax_params = collections.OrderedDict() 114 | 115 | for k in jax_params: 116 | if "UpSample" not in k: 117 | params = jax_params[k] 118 | 119 | if ("ConvTranspose" in k) and ("bias" not in k): 120 | params = params.transpose(0, 1, 3, 2) 121 | 122 | split_names = k.split("_") 123 | modified_param_name = ( 124 | "_".join(split_names[0:-1]) + "/" + split_names[-1] + ":0" 125 | ) 126 | 127 | if "layernorm" in modified_param_name.lower(): 128 | if "scale" in modified_param_name: 129 | modified_param_name = modified_param_name.replace("scale", "gamma") 130 | elif "bias" in modified_param_name: 131 | modified_param_name = modified_param_name.replace("bias", "beta") 132 | 133 | modified_jax_params.update({modified_param_name: params}) 134 | 135 | return modified_jax_params 136 | 137 | 138 | def port_jax_params(configs: dict, ckpt_path: str) -> Tuple[dict, tf.keras.Model]: 139 | # Initialize TF Model. 140 | print("Initializing model.") 141 | tf_model = Model(**configs) 142 | 143 | # Obtain a mapping of the TF variable names and their values. 144 | tf_model_variables = tf_model.variables 145 | tf_model_variables_dict = {} 146 | for v in tf_model_variables: 147 | tf_model_variables_dict[v.name] = v 148 | 149 | # Obtain the JAX pre-trained variables. 150 | jax_params = get_params(ckpt_path) 151 | [flat_jax_dict] = pd.json_normalize(jax_params, sep="_").to_dict(orient="records") 152 | 153 | # Amend the JAX variables to match the names of the TF variables. 154 | modified_jax_params = modify_jax_params(flat_jax_dict) 155 | modified_jax_params.update(modify_upsample(jax_params)) 156 | 157 | # Porting. 158 | tf_weights = [] 159 | i = 0 160 | 161 | for k in modified_jax_params: 162 | param = modified_jax_params[k] 163 | tf_weights.append((tf_model_variables_dict[k], param)) 164 | i += 1 165 | 166 | assert i == len(modified_jax_params) == len(tf_model_variables_dict) 167 | 168 | tf.keras.backend.batch_set_value(tf_weights) 169 | 170 | return modified_jax_params, tf_model 171 | 172 | 173 | def main(args): 174 | task = args.task 175 | task_from_ckpt = args.ckpt_path.split("/")[-3] 176 | 177 | assert task == task_from_ckpt, "Provided task and provided checkpoints differ." 178 | f" Task provided: {task}, task dervived from checkpoints: {task_from_ckpt}." 179 | 180 | # From https://github.com/google-research/maxim/blob/main/maxim/run_eval.py#L55 181 | variant = _MODEL_VARIANT_DICT[task] 182 | configs = MAXIM_CONFIGS.get(variant) 183 | configs.update( 184 | { 185 | "variant": variant, 186 | "dropout_rate": 0.0, 187 | "num_outputs": 3, 188 | "use_bias": True, 189 | "num_supervision_scales": 3, 190 | } 191 | ) 192 | 193 | _, tf_model = port_jax_params(configs, args.ckpt_path) 194 | print("Model porting successful.") 195 | 196 | dataset_name = args.ckpt_path.split("/")[-2].lower() 197 | tf_params_path = f"{variant}_{task.lower()}_{dataset_name}.h5" 198 | 199 | tf_model.save_weights(tf_params_path) 200 | print(f"Model params serialized to {tf_params_path}.") 201 | saved_model_path = tf_params_path.replace(".h5", "") 202 | push_to_hub_keras(tf_model, repo_path_or_name=f"sayakpaul/{saved_model_path}") 203 | print("Model pushed to Hugging Face Hub.") 204 | 205 | 206 | def parse_args(): 207 | parser = argparse.ArgumentParser( 208 | description="Conversion of the JAX pre-trained MAXIM weights to TensorFlow." 209 | ) 210 | parser.add_argument( 211 | "-t", 212 | "--task", 213 | default="Denoising", 214 | type=str, 215 | choices=[ 216 | "Denoising", 217 | "Deblurring", 218 | "Deraining", 219 | "Dehazing", 220 | "Enhancement", 221 | ], 222 | help="Name of the task on which the corresponding checkpoints were derived.", 223 | ) 224 | parser.add_argument( 225 | "-c", 226 | "--ckpt_path", 227 | default="gs://gresearch/maxim/ckpt/Denoising/SIDD/checkpoint.npz", 228 | type=str, 229 | help="Checkpoint to port.", 230 | ) 231 | return parser.parse_args() 232 | 233 | 234 | if __name__ == "__main__": 235 | args = parse_args() 236 | main(args) 237 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # MAXIM in TensorFlow 2 | 3 | [![HugginFace badge](https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-yellow.svg)](https://huggingface.co/spaces/sayakpaul/maxim-spaces) 4 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sayakpaul/maxim-tf/blob/main/notebooks/inference-dynamic-resize.ipynb) [![TensorFlow 2.10](https://img.shields.io/badge/TensorFlow-2.10-FF6F00?logo=tensorflow)](https://github.com/tensorflow/tensorflow/releases/tag/v2.8.0) 5 | [![Models on TF-Hub](https://img.shields.io/badge/TF--Hub-Models%20on%20TF--Hub-orange)](https://tfhub.dev/sayakpaul/collections/maxim/1) 6 | [![HugginFace badge](https://img.shields.io/badge/🤗%20Hugging%20Face-Hub-yellow.svg)](https://huggingface.co/models?pipeline_tag=image-to-image&sort=downloads&search=maxim) 7 | 8 | Implementation of MAXIM [1] in TensorFlow. This project received the [#TFCommunitySpotlight Award](https://twitter.com/TensorFlow/status/1611469033714470919?s=20&t=a5LMpYXrPg6E0WGudsYezw). 9 | 10 | MAXIM introduces a backbone that can tackle image denoising, dehazing, deblurring, deraining, and enhancement. 11 | 12 |
13 | 14 | Taken from the MAXIM paper 15 |
16 | 17 | The weights of different MAXIM variants are in JAX and they're available in [2]. 18 | 19 | You can find all the TensorFlow MAXIM models [here](https://tfhub.dev/sayakpaul/collections/maxim/1) on TensorFlow Hub as 20 | well as on [Hugging Face Hub](https://huggingface.co/models?pipeline_tag=image-to-image&sort=downloads&search=maxim). 21 | 22 | You can try out the models on Hugging Face Spaces: 23 | 24 | * [Denoising](https://huggingface.co/spaces/sayakpaul/sidd-denoising-maxim) 25 | * [Low-light enhancement](https://huggingface.co/spaces/sayakpaul/lol-enhancement-maxim) 26 | * [Image retouching](https://huggingface.co/spaces/sayakpaul/fivek-retouching-maxim) 27 | * [Dehazing indoors](https://huggingface.co/spaces/sayakpaul/sots-indoor-dehazing-maxim) 28 | * [Dehazing outdoors](https://huggingface.co/spaces/sayakpaul/sots-outdoor-dehazing-maxim) 29 | * [Image deraining](https://huggingface.co/spaces/sayakpaul/rain13k-deraining-maxim) 30 | * [Image deblurring](https://huggingface.co/spaces/sayakpaul/gopro-deblurring-maxim) 31 | 32 | If you prefer Colab Notebooks, then you can check them out [here](https://github.com/sayakpaul/maxim-tf/tree/main/notebooks). 33 | ## Model conversion to TensorFlow from JAX 34 | 35 | Blocks and layers related to MAXIM are implemented in the `maxim` directory. 36 | 37 | `convert_to_tf.py` script is leveraged to initialize a particular MAXIM model variant and a pre-trained checkpoint and then run the conversion to TensorFlow. Refer to the usage section of the script to know more. 38 | 39 | This script serializes the model weights in `.h5` as as well pushes the `SavedModel` to Hugging Face Hub. For the latter, you need to authenticate yourself if not already done (`huggingface-cli login`). 40 | 41 | This TensorFlow implementation is in close alignment with [2]. The author of this repository has reused some code blocks from [2] (with credits) to do. 42 | 43 | ## Results and model variants 44 | 45 | A comprehensive table is available [here](https://github.com/google-research/maxim#results-and-pre-trained-models). The author of this repository validated the results with the converted models qualitatively. 46 | 47 |
48 | 49 |
50 | 51 | ## Inference with the provided sample images 52 | 53 | You can run the `run_eval.py` script for this purpose. 54 | 55 |
56 | Image Denoising (click to expand) 57 | 58 | ``` 59 | python3 maxim/run_eval.py --task Denoising --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-3_denoising_sidd/1/uncompressed \ 60 | --input_dir images/Denoising --output_dir images/Results --has_target=False --dynamic_resize=True 61 | ``` 62 |
63 | 64 |
65 | Image Deblurring (click to expand) 66 | 67 | ``` 68 | python3 maxim/run_eval.py --task Deblurring --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-3_deblurring_gopro/1/uncompressed \ 69 | --input_dir images/Deblurring --output_dir images/Results --has_target=False --dynamic_resize=True 70 | ``` 71 |
72 | 73 |
74 | Image Deraining (click to expand) 75 | 76 | Rain streak: 77 | ``` 78 | python3 maxim/run_eval.py --task Deraining --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_deraining_rain13k/1/uncompressed \ 79 | --input_dir images/Deraining --output_dir images/Results --has_target=False --dynamic_resize=True 80 | ``` 81 | 82 | Rain drop: 83 | ``` 84 | python3 maxim/run_eval.py --task Deraining --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_deraining_raindrop/1/uncompressed \ 85 | --input_dir images/Deraining --output_dir images/Results --has_target=False --dynamic_resize=True 86 | ``` 87 |
88 | 89 |
90 | Image Dehazing (click to expand) 91 | 92 | Indoor: 93 | ``` 94 | python3 maxim/run_eval.py --task Dehazing --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_dehazing_sots-indoor/1/uncompressed \ 95 | --input_dir images/Dehazing --output_dir images/Results --has_target=False --dynamic_resize=True 96 | ``` 97 | 98 | Outdoor: 99 | ``` 100 | python3 maxim/run_eval.py --task Dehazing --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_dehazing_sots-outdoor/1/uncompressed \ 101 | --input_dir images/Dehazing --output_dir images/Results --has_target=False --dynamic_resize=True 102 | ``` 103 |
104 | 105 |
106 | Image Enhancement (click to expand) 107 | 108 | Low-light enhancement: 109 | ``` 110 | python3 maxim/run_eval.py --task Enhancement --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_enhancement_lol/1/uncompressed \ 111 | --input_dir images/Enhancement --output_dir images/Results --has_target=False --dynamic_resize=True 112 | ``` 113 | 114 | Retouching: 115 | ``` 116 | python3 maxim/run_eval.py --task Enhancement --ckpt_path gs://tfhub-modules/sayakpaul/maxim_s-2_enhancement_fivek/1/uncompressed \ 117 | --input_dir images/Enhancement --output_dir images/Results --has_target=False --dynamic_resize=True 118 | ``` 119 |
120 |
121 | 122 | **Notes**: 123 | 124 | * The `run_eval.py` script is heavily inspired by the [original one](https://github.com/google-research/maxim/blob/main/maxim/run_eval.py). 125 | * You can set `dynamic_resize` to False to obtain faster latency compromising the prediction quality. 126 | 127 | 128 | ## XLA support 129 | 130 | The models are XLA-supported. It can drammatically reduce the latency. Refer to the `benchmark_xla.py` script for more. 131 | 132 | ## Known limitations 133 | 134 | These are some of the known limitations of the current implementation. These are all 135 | open for contributions. 136 | 137 | ### Supporting arbitrary image resolutions 138 | 139 | MAXIM supports arbitrary image resolutions. However, the available TensorFlow models were exported with `(256, 256, 3)` resolution. So, a crude form of resizing is done on the input images to perform inference with the available models. This impacts the results quite a bit. This issue is discussed in more details [here](https://github.com/sayakpaul/maxim-tf/issues/11). [Some work](https://github.com/sayakpaul/maxim-tf/pull/20) has been started to fix this behaviour (without ETA). I am thankful to [Amy Roberts](https://uk.linkedin.com/in/amy-roberts-70903a6a) from Hugging Face for guiding me in the right direction. 140 | 141 | But these models can be extended to support arbitrary resolution. Refer to [this notebook](https://colab.research.google.com/github/sayakpaul/maxim-tf/blob/main/notebooks/inference-dynamic-resize.ipynb) for more details. Specifically, for a given task and an image, a new version of the model is instantiated and the weights of the available model are copied into the new model instance. This is a time-consuming process and isn't very efficient. 142 | 143 | #### Changes to achieve arbitrary image resolution on TF 144 | 145 | - Substitute einops calls for pure TF operations: Einops operations are not intended operate on data-dependent (unknown) dimensionality [https://github.com/data-apis/array-api/issues/494](https://github.com/data-apis/array-api/issues/494). Thus, it was necessary to re-write BlockImages and UnblockImages as full TF ops. For convenience, we separate BlockImages into TFBlockImages and TFBlockImagesByGrid. We also rewrote UnblockImages as TFUnblockImages. 146 | - Make [dim_u](https://github.com/sayakpaul/maxim-tf/pull/24/files#diff-8b281bcfc137b53489e1b19b29735462d5deac19b8c2c2f82cf0383680908063R121) and [dim_v](https://github.com/sayakpaul/maxim-tf/pull/24/files#diff-8b281bcfc137b53489e1b19b29735462d5deac19b8c2c2f82cf0383680908063R130) parameters independent of the input image size. This can be done by computing dim_u and dim_v from the provided grid_size and/or block_size. 147 | - Change resizing layers so as to receive a ratio independent of the image size. It was important to use the float ratios to compute the final image size, just then converting back to int, to avoid loss of information. 148 | 149 | ### Output mismatches 150 | 151 | The outputs of the TF and JAX models vary slightly. This is because of the differences in the implementation of different layers (resizing layer mainly). Even though the differences in the outputs of individual blocks of TF and JAX models are small, they add up, in the end, to be larger than one might expect. 152 | 153 | With all that said, the qualitative performance doesn't seem to be disturbed at all. 154 | 155 | ## Call for contributions 156 | 157 | - [ ] Add a minimal training notebook. 158 | - [ ] Fix any of the known limitations stated above 159 | 160 | ## Acknowledgements 161 | 162 | * ML Developer Programs' team at Google for providing Google Cloud credits. 163 | * [Gustavo Martins](https://twitter.com/gusthema?lang=en) from Google for initial discussions and reviews of the codebase. 164 | * [Amy Roberts](https://uk.linkedin.com/in/amy-roberts-70903a6a) from Hugging Face for guiding me in the right direction for handling arbitrary input shapes. 165 | 166 | ## References 167 | 168 | [1] MAXIM paper: https://arxiv.org/abs/2201.02973 169 | 170 | [2] MAXIM official GitHub: https://github.com/google-research/maxim 171 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /notebooks/inference.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "id": "view-in-github", 7 | "colab_type": "text" 8 | }, 9 | "source": [ 10 | "\"Open" 11 | ] 12 | }, 13 | { 14 | "cell_type": "markdown", 15 | "metadata": { 16 | "id": "Wxgb-sWfpb49" 17 | }, 18 | "source": [ 19 | "## Introduction\n", 20 | "\n", 21 | "This notebook shows how to run inference with the [MAXIM family of models](https://github.com/google-research/maxim) from [TensorFlow Hub](https://tfhub.dev/sayakpaul/collections/maxim/1). MAXIM family of models share the same backbone for performing: denoising, dehazing, deblurring, deraining, and enhancement. You can know more about the public MAXIM models from [here](https://github.com/google-research/maxim#results-and-pre-trained-models)." 22 | ] 23 | }, 24 | { 25 | "cell_type": "markdown", 26 | "metadata": { 27 | "id": "Zlp7twW3tB2n" 28 | }, 29 | "source": [ 30 | "## Select a checkpoint" 31 | ] 32 | }, 33 | { 34 | "cell_type": "code", 35 | "execution_count": null, 36 | "metadata": { 37 | "id": "E-n8jA4Gojv2", 38 | "cellView": "form" 39 | }, 40 | "outputs": [], 41 | "source": [ 42 | "task = \"Dehazing_Indoor\" # @param [\"Denoising\", \"Dehazing_Indoor\", \"Dehazing_Outdoor\", \"Deblurring\", \"Deraining\", \"Enhancement\", \"Retouching\"]\n", 43 | "\n", 44 | "model_handle_map = {\n", 45 | " \"Denoising\": [\n", 46 | " \"https://tfhub.dev/sayakpaul/maxim_s-3_denoising_sidd/1\",\n", 47 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Denoising/input/0003_30.png\",\n", 48 | " ],\n", 49 | " \"Dehazing_Indoor\": [\n", 50 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_dehazing_sots-indoor/1\",\n", 51 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Dehazing/input/0003_0.8_0.2.png\",\n", 52 | " ],\n", 53 | " \"Dehazing_Outdoor\": [\n", 54 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_dehazing_sots-outdoor/1\",\n", 55 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Dehazing/input/1444_10.png\",\n", 56 | " ],\n", 57 | " \"Deblurring\": [\n", 58 | " \"https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_gopro/1\",\n", 59 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Deblurring/input/1fromGOPR0950.png\",\n", 60 | " ],\n", 61 | " \"Deraining\": [\n", 62 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_deraining_raindrop/1\",\n", 63 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Deraining/input/15.png\",\n", 64 | " ],\n", 65 | " \"Enhancement\": [\n", 66 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_enhancement_lol/1\",\n", 67 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Enhancement/input/a4541-DSC_0040-2.png\",\n", 68 | " ],\n", 69 | " \"Retouching\": [\n", 70 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_enhancement_fivek/1\",\n", 71 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Enhancement/input/a4541-DSC_0040-2.png\",\n", 72 | " ],\n", 73 | "}\n", 74 | "\n", 75 | "model_handle = model_handle_map[task]\n", 76 | "ckpt = model_handle[0]\n", 77 | "print(f\"TF-Hub handle: {ckpt}.\")" 78 | ] 79 | }, 80 | { 81 | "cell_type": "markdown", 82 | "metadata": { 83 | "id": "3t6c3Z1Pz6UT" 84 | }, 85 | "source": [ 86 | "For deblurring, there are other checkpoints too:\n", 87 | "\n", 88 | "* https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_realblur_r/1\n", 89 | "* https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_realblur_j/1\n", 90 | "* https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_reds/1" 91 | ] 92 | }, 93 | { 94 | "cell_type": "markdown", 95 | "metadata": { 96 | "id": "SNlQ2HzrtJOU" 97 | }, 98 | "source": [ 99 | "## Imports" 100 | ] 101 | }, 102 | { 103 | "cell_type": "code", 104 | "execution_count": null, 105 | "metadata": { 106 | "id": "IEhpgokqtKFz" 107 | }, 108 | "outputs": [], 109 | "source": [ 110 | "import tensorflow as tf\n", 111 | "import tensorflow_hub as hub\n", 112 | "\n", 113 | "import matplotlib.pyplot as plt\n", 114 | "\n", 115 | "from PIL import Image\n", 116 | "import numpy as np" 117 | ] 118 | }, 119 | { 120 | "cell_type": "markdown", 121 | "metadata": { 122 | "id": "UTdrCvUltkCn" 123 | }, 124 | "source": [ 125 | "## Fetch the input image based on the task" 126 | ] 127 | }, 128 | { 129 | "cell_type": "code", 130 | "execution_count": null, 131 | "metadata": { 132 | "id": "Bn90H1rltRcM" 133 | }, 134 | "outputs": [], 135 | "source": [ 136 | "image_url = model_handle[1]\n", 137 | "image_path = tf.keras.utils.get_file(origin=image_url)\n", 138 | "Image.open(image_path)" 139 | ] 140 | }, 141 | { 142 | "cell_type": "markdown", 143 | "metadata": { 144 | "id": "UsXRY1kvum4O" 145 | }, 146 | "source": [ 147 | "## Preprocessing utilities\n", 148 | "\n", 149 | "Based on [this official script](https://github.com/google-research/maxim/blob/main/maxim/run_eval.py)." 150 | ] 151 | }, 152 | { 153 | "cell_type": "code", 154 | "execution_count": null, 155 | "metadata": { 156 | "id": "ZGmgEtfLt7S5" 157 | }, 158 | "outputs": [], 159 | "source": [ 160 | "# Since the model was not initialized to take variable-length sizes (None, None, 3),\n", 161 | "# we need to be careful about how we are resizing the images.\n", 162 | "# From https://www.tensorflow.org/lite/examples/style_transfer/overview#pre-process_the_inputs\n", 163 | "def resize_image(image, target_dim):\n", 164 | " # Resize the image so that the shorter dimension becomes `target_dim`.\n", 165 | " shape = tf.cast(tf.shape(image)[1:-1], tf.float32)\n", 166 | " short_dim = min(shape)\n", 167 | " scale = target_dim / short_dim\n", 168 | " new_shape = tf.cast(shape * scale, tf.int32)\n", 169 | " image = tf.image.resize(image, new_shape)\n", 170 | "\n", 171 | " # Central crop the image.\n", 172 | " image = tf.image.resize_with_crop_or_pad(image, target_dim, target_dim)\n", 173 | "\n", 174 | " return image\n", 175 | "\n", 176 | "\n", 177 | "def process_image(image_path, target_dim=256):\n", 178 | " input_img = np.asarray(Image.open(image_path).convert(\"RGB\"), np.float32) / 255.0\n", 179 | " input_img = tf.expand_dims(input_img, axis=0)\n", 180 | " input_img = resize_image(input_img, target_dim)\n", 181 | " return input_img" 182 | ] 183 | }, 184 | { 185 | "cell_type": "markdown", 186 | "metadata": { 187 | "id": "FxsxGhvKvDrF" 188 | }, 189 | "source": [ 190 | "This notebook infers on fixed-shape images. However, MAXIM can handle images of any resolution. The current implementation in TensorFlow can achieve this with a bit of hacking. Please refer to [this notebook](https://github.com/sayakpaul/maxim-tf/blob/main/notebooks/inference-dynamic-resize.ipynb) if you want the model to infer on dynamic shapes. " 191 | ] 192 | }, 193 | { 194 | "cell_type": "markdown", 195 | "metadata": { 196 | "id": "5T20A1xLvcLq" 197 | }, 198 | "source": [ 199 | "## Run predictions" 200 | ] 201 | }, 202 | { 203 | "cell_type": "code", 204 | "execution_count": null, 205 | "metadata": { 206 | "id": "hOxESHl2vdRE" 207 | }, 208 | "outputs": [], 209 | "source": [ 210 | "def get_model(model_url: str, input_resolution: tuple) -> tf.keras.Model:\n", 211 | " inputs = tf.keras.Input((*input_resolution, 3))\n", 212 | " hub_module = hub.KerasLayer(model_url)\n", 213 | "\n", 214 | " outputs = hub_module(inputs)\n", 215 | "\n", 216 | " return tf.keras.Model(inputs, outputs)\n", 217 | "\n", 218 | "\n", 219 | "# Based on https://github.com/google-research/maxim/blob/main/maxim/run_eval.py\n", 220 | "def infer(image_path: str, model: tf.keras.Model, input_resolution=(256, 256)):\n", 221 | " preprocessed_image = process_image(image_path, input_resolution[0])\n", 222 | "\n", 223 | " preds = model.predict(preprocessed_image)\n", 224 | " if isinstance(preds, list):\n", 225 | " preds = preds[-1]\n", 226 | " if isinstance(preds, list):\n", 227 | " preds = preds[-1]\n", 228 | "\n", 229 | " preds = np.array(preds[0], np.float32)\n", 230 | " final_pred_image = np.array((np.clip(preds, 0.0, 1.0)).astype(np.float32))\n", 231 | " return final_pred_image" 232 | ] 233 | }, 234 | { 235 | "cell_type": "code", 236 | "execution_count": null, 237 | "metadata": { 238 | "id": "1Fr-rYLpwab6" 239 | }, 240 | "outputs": [], 241 | "source": [ 242 | "input_resolution = (256, 256)\n", 243 | "\n", 244 | "model = get_model(ckpt, input_resolution)\n", 245 | "\n", 246 | "final_pred_image = infer(image_path, model, input_resolution)" 247 | ] 248 | }, 249 | { 250 | "cell_type": "markdown", 251 | "metadata": { 252 | "id": "GTq1J42tw67G" 253 | }, 254 | "source": [ 255 | "## Visualize results" 256 | ] 257 | }, 258 | { 259 | "cell_type": "code", 260 | "execution_count": null, 261 | "metadata": { 262 | "id": "ECGdFWQBw8E2" 263 | }, 264 | "outputs": [], 265 | "source": [ 266 | "# Based on https://www.tensorflow.org/lite/examples/style_transfer/overview#visualize_the_inputs\n", 267 | "def imshow(image, title=None):\n", 268 | " if len(image.shape) > 3:\n", 269 | " image = tf.squeeze(image, axis=0)\n", 270 | "\n", 271 | " plt.imshow(image)\n", 272 | " if title:\n", 273 | " plt.title(title)\n", 274 | "\n", 275 | "\n", 276 | "plt.figure(figsize=(15, 15))\n", 277 | "\n", 278 | "plt.subplot(1, 2, 1)\n", 279 | "input_image = np.asarray(Image.open(image_path).convert(\"RGB\"), np.float32) / 255.0\n", 280 | "imshow(input_image, \"Input Image\")\n", 281 | "\n", 282 | "plt.subplot(1, 2, 2)\n", 283 | "imshow(final_pred_image, \"Predicted Image\")" 284 | ] 285 | } 286 | ], 287 | "metadata": { 288 | "accelerator": "GPU", 289 | "colab": { 290 | "provenance": [], 291 | "include_colab_link": true 292 | }, 293 | "kernelspec": { 294 | "display_name": "Python 3 (ipykernel)", 295 | "language": "python", 296 | "name": "python3" 297 | }, 298 | "language_info": { 299 | "codemirror_mode": { 300 | "name": "ipython", 301 | "version": 3 302 | }, 303 | "file_extension": ".py", 304 | "mimetype": "text/x-python", 305 | "name": "python", 306 | "nbconvert_exporter": "python", 307 | "pygments_lexer": "ipython3", 308 | "version": "3.8.2" 309 | } 310 | }, 311 | "nbformat": 4, 312 | "nbformat_minor": 0 313 | } -------------------------------------------------------------------------------- /notebooks/inference-dynamic-resize.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "id": "view-in-github", 7 | "colab_type": "text" 8 | }, 9 | "source": [ 10 | "\"Open" 11 | ] 12 | }, 13 | { 14 | "cell_type": "markdown", 15 | "metadata": { 16 | "id": "Wxgb-sWfpb49" 17 | }, 18 | "source": [ 19 | "## Introduction\n", 20 | "\n", 21 | "This notebook shows how to run inference with the [MAXIM family of models](https://github.com/google-research/maxim) from [TensorFlow Hub](https://tfhub.dev/sayakpaul/collections/maxim/1). MAXIM family of models share the same backbone for performing: denoising, dehazing, deblurring, deraining, and enhancement. You can know more about the public MAXIM models from [here](https://github.com/google-research/maxim#results-and-pre-trained-models).\n", 22 | "\n", 23 | "This notebook allows you to run dynamic shaped images unlike [this one](https://github.com/sayakpaul/maxim-tf/blob/main/notebooks/inference.ipynb)." 24 | ] 25 | }, 26 | { 27 | "cell_type": "markdown", 28 | "metadata": { 29 | "id": "Zlp7twW3tB2n" 30 | }, 31 | "source": [ 32 | "## Select a checkpoint" 33 | ] 34 | }, 35 | { 36 | "cell_type": "code", 37 | "execution_count": null, 38 | "metadata": { 39 | "id": "E-n8jA4Gojv2" 40 | }, 41 | "outputs": [], 42 | "source": [ 43 | "task = \"Deblurring\" # @param [\"Denoising\", \"Dehazing_Indoor\", \"Dehazing_Outdoor\", \"Deblurring\", \"Deraining\", \"Enhancement\", \"Retouching\"]\n", 44 | "\n", 45 | "model_handle_map = {\n", 46 | " \"Denoising\": [\n", 47 | " \"https://tfhub.dev/sayakpaul/maxim_s-3_denoising_sidd/1\",\n", 48 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Denoising/input/0003_30.png\",\n", 49 | " ],\n", 50 | " \"Dehazing_Indoor\": [\n", 51 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_dehazing_sots-indoor/1\",\n", 52 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Dehazing/input/0003_0.8_0.2.png\",\n", 53 | " ],\n", 54 | " \"Dehazing_Outdoor\": [\n", 55 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_dehazing_sots-outdoor/1\",\n", 56 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Dehazing/input/1444_10.png\",\n", 57 | " ],\n", 58 | " \"Deblurring\": [\n", 59 | " \"https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_gopro/1\",\n", 60 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Deblurring/input/1fromGOPR0950.png\",\n", 61 | " ],\n", 62 | " \"Deraining\": [\n", 63 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_deraining_raindrop/1\",\n", 64 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Deraining/input/15.png\",\n", 65 | " ],\n", 66 | " \"Enhancement\": [\n", 67 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_enhancement_lol/1\",\n", 68 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Enhancement/input/a4541-DSC_0040-2.png\",\n", 69 | " ],\n", 70 | " \"Retouching\": [\n", 71 | " \"https://tfhub.dev/sayakpaul/maxim_s-2_enhancement_fivek/1\",\n", 72 | " \"https://github.com/google-research/maxim/raw/main/maxim/images/Enhancement/input/a4541-DSC_0040-2.png\",\n", 73 | " ],\n", 74 | "}\n", 75 | "\n", 76 | "model_handle = model_handle_map[task]\n", 77 | "ckpt = model_handle[0]\n", 78 | "print(f\"TF-Hub handle: {ckpt}.\")" 79 | ] 80 | }, 81 | { 82 | "cell_type": "markdown", 83 | "metadata": { 84 | "id": "3t6c3Z1Pz6UT" 85 | }, 86 | "source": [ 87 | "For deblurring, there are other checkpoints too:\n", 88 | "\n", 89 | "- https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_realblur_r/1\n", 90 | "- https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_realblur_j/1\n", 91 | "- https://tfhub.dev/sayakpaul/maxim_s-3_deblurring_reds/1\n" 92 | ] 93 | }, 94 | { 95 | "cell_type": "markdown", 96 | "metadata": { 97 | "id": "SNlQ2HzrtJOU" 98 | }, 99 | "source": [ 100 | "## Imports" 101 | ] 102 | }, 103 | { 104 | "cell_type": "code", 105 | "execution_count": null, 106 | "metadata": { 107 | "id": "IEhpgokqtKFz" 108 | }, 109 | "outputs": [], 110 | "source": [ 111 | "import tensorflow as tf\n", 112 | "import tensorflow_hub as hub\n", 113 | "\n", 114 | "import matplotlib.pyplot as plt\n", 115 | "\n", 116 | "from PIL import Image\n", 117 | "import numpy as np" 118 | ] 119 | }, 120 | { 121 | "cell_type": "code", 122 | "execution_count": null, 123 | "metadata": { 124 | "id": "us19UdxvdE18" 125 | }, 126 | "outputs": [], 127 | "source": [ 128 | "import sys\n", 129 | "\n", 130 | "sys.path.append(\"..\")\n", 131 | "\n", 132 | "from create_maxim_model import Model\n", 133 | "from maxim.configs import MAXIM_CONFIGS" 134 | ] 135 | }, 136 | { 137 | "cell_type": "markdown", 138 | "metadata": { 139 | "id": "YEtzvxy5dE18" 140 | }, 141 | "source": [ 142 | "TODO: When the repository is public, clone it and use accordingly." 143 | ] 144 | }, 145 | { 146 | "cell_type": "markdown", 147 | "metadata": { 148 | "id": "UTdrCvUltkCn" 149 | }, 150 | "source": [ 151 | "## Fetch the input image based on the task" 152 | ] 153 | }, 154 | { 155 | "cell_type": "code", 156 | "execution_count": null, 157 | "metadata": { 158 | "id": "Bn90H1rltRcM" 159 | }, 160 | "outputs": [], 161 | "source": [ 162 | "image_url = model_handle[1]\n", 163 | "image_path = tf.keras.utils.get_file(origin=image_url)\n", 164 | "Image.open(image_path)" 165 | ] 166 | }, 167 | { 168 | "cell_type": "markdown", 169 | "metadata": { 170 | "id": "qs1FCtINdE19" 171 | }, 172 | "source": [ 173 | "## Load the model" 174 | ] 175 | }, 176 | { 177 | "cell_type": "code", 178 | "execution_count": null, 179 | "metadata": { 180 | "id": "7dzkVpOQdE1-" 181 | }, 182 | "outputs": [], 183 | "source": [ 184 | "_MODEL = tf.keras.models.load_model(ckpt)" 185 | ] 186 | }, 187 | { 188 | "cell_type": "markdown", 189 | "metadata": { 190 | "id": "UsXRY1kvum4O" 191 | }, 192 | "source": [ 193 | "## Preprocessing utilities\n", 194 | "\n", 195 | "Based on [this official script](https://github.com/google-research/maxim/blob/main/maxim/run_eval.py)." 196 | ] 197 | }, 198 | { 199 | "cell_type": "code", 200 | "execution_count": null, 201 | "metadata": { 202 | "id": "ZGmgEtfLt7S5" 203 | }, 204 | "outputs": [], 205 | "source": [ 206 | "def mod_padding_symmetric(image, factor=64):\n", 207 | " \"\"\"Padding the image to be divided by factor.\"\"\"\n", 208 | " height, width = image.shape[0], image.shape[1]\n", 209 | " height_pad, width_pad = ((height + factor) // factor) * factor, (\n", 210 | " (width + factor) // factor\n", 211 | " ) * factor\n", 212 | " padh = height_pad - height if height % factor != 0 else 0\n", 213 | " padw = width_pad - width if width % factor != 0 else 0\n", 214 | " image = tf.pad(\n", 215 | " image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)], mode=\"REFLECT\"\n", 216 | " )\n", 217 | " return image\n", 218 | "\n", 219 | "\n", 220 | "def make_shape_even(image):\n", 221 | " \"\"\"Pad the image to have even shapes.\"\"\"\n", 222 | " height, width = image.shape[0], image.shape[1]\n", 223 | " padh = 1 if height % 2 != 0 else 0\n", 224 | " padw = 1 if width % 2 != 0 else 0\n", 225 | " image = tf.pad(image, [(0, padh), (0, padw), (0, 0)], mode=\"REFLECT\")\n", 226 | " return image\n", 227 | "\n", 228 | "\n", 229 | "def process_image(image: Image):\n", 230 | " input_img = np.asarray(image) / 255.0\n", 231 | " height, width = input_img.shape[0], input_img.shape[1]\n", 232 | "\n", 233 | " # Padding images to have even shapes\n", 234 | " input_img = make_shape_even(input_img)\n", 235 | " height_even, width_even = input_img.shape[0], input_img.shape[1]\n", 236 | "\n", 237 | " # padding images to be multiplies of 64\n", 238 | " input_img = mod_padding_symmetric(input_img, factor=64)\n", 239 | " input_img = tf.expand_dims(input_img, axis=0)\n", 240 | " return input_img, height, width, height_even, width_even\n", 241 | "\n", 242 | "\n", 243 | "def init_new_model(input_img):\n", 244 | " variant = ckpt.split(\"/\")[-1].split(\"_\")[0]\n", 245 | " configs = MAXIM_CONFIGS.get(variant)\n", 246 | " configs.update(\n", 247 | " {\n", 248 | " \"variant\": \"S-2\",\n", 249 | " \"dropout_rate\": 0.0,\n", 250 | " \"num_outputs\": 3,\n", 251 | " \"use_bias\": True,\n", 252 | " \"num_supervision_scales\": 3,\n", 253 | " }\n", 254 | " ) # From https://github.com/google-research/maxim/blob/main/maxim/run_eval.py#L45-#L61\n", 255 | " configs.update({\"input_resolution\": (input_img.shape[1], input_img.shape[2])})\n", 256 | " new_model = Model(**configs)\n", 257 | " new_model.set_weights(_MODEL.get_weights())\n", 258 | " return new_model" 259 | ] 260 | }, 261 | { 262 | "cell_type": "markdown", 263 | "metadata": { 264 | "id": "lVOdyjTudE1_" 265 | }, 266 | "source": [ 267 | "To make the model operate on images of arbitrary shapes here's what we're doing:\n", 268 | "\n", 269 | "* Loading the initial pre-trained model into `_MODEL`.\n", 270 | "* Initializing a separate instance of MAXIM based on the configs and spatial resolutions of the input image.\n", 271 | "* Populating the params of this newly initialized model with that of `_MODEL`. \n", 272 | "\n", 273 | "All of it is handled in `init_new_model()`. " 274 | ] 275 | }, 276 | { 277 | "cell_type": "markdown", 278 | "metadata": { 279 | "id": "5T20A1xLvcLq" 280 | }, 281 | "source": [ 282 | "## Run predictions" 283 | ] 284 | }, 285 | { 286 | "cell_type": "code", 287 | "execution_count": null, 288 | "metadata": { 289 | "id": "hOxESHl2vdRE" 290 | }, 291 | "outputs": [], 292 | "source": [ 293 | "# Based on https://github.com/google-research/maxim/blob/main/maxim/run_eval.py\n", 294 | "def infer(image_path: str):\n", 295 | " image = Image.open(image_path).convert(\"RGB\")\n", 296 | " preprocessed_image, height, width, height_even, width_even = process_image(image)\n", 297 | " new_model = init_new_model(preprocessed_image)\n", 298 | "\n", 299 | " preds = new_model.predict(preprocessed_image)\n", 300 | " if isinstance(preds, list):\n", 301 | " preds = preds[-1]\n", 302 | " if isinstance(preds, list):\n", 303 | " preds = preds[-1]\n", 304 | "\n", 305 | " preds = np.array(preds[0], np.float32)\n", 306 | "\n", 307 | " new_height, new_width = preds.shape[0], preds.shape[1]\n", 308 | " h_start = new_height // 2 - height_even // 2\n", 309 | " h_end = h_start + height\n", 310 | " w_start = new_width // 2 - width_even // 2\n", 311 | " w_end = w_start + width\n", 312 | " preds = preds[h_start:h_end, w_start:w_end, :]\n", 313 | "\n", 314 | " return np.array(np.clip(preds, 0.0, 1.0))" 315 | ] 316 | }, 317 | { 318 | "cell_type": "code", 319 | "execution_count": null, 320 | "metadata": { 321 | "id": "1Fr-rYLpwab6" 322 | }, 323 | "outputs": [], 324 | "source": [ 325 | "final_pred_image = infer(image_path)" 326 | ] 327 | }, 328 | { 329 | "cell_type": "markdown", 330 | "metadata": { 331 | "id": "GTq1J42tw67G" 332 | }, 333 | "source": [ 334 | "## Visualize results" 335 | ] 336 | }, 337 | { 338 | "cell_type": "code", 339 | "execution_count": null, 340 | "metadata": { 341 | "id": "ECGdFWQBw8E2" 342 | }, 343 | "outputs": [], 344 | "source": [ 345 | "# Based on https://www.tensorflow.org/lite/examples/style_transfer/overview#visualize_the_inputs\n", 346 | "def imshow(image, title=None):\n", 347 | " if len(image.shape) > 3:\n", 348 | " image = tf.squeeze(image, axis=0)\n", 349 | "\n", 350 | " plt.imshow(image)\n", 351 | " if title:\n", 352 | " plt.title(title)\n", 353 | "\n", 354 | "\n", 355 | "plt.figure(figsize=(15, 15))\n", 356 | "\n", 357 | "plt.subplot(1, 2, 1)\n", 358 | "input_image = np.asarray(Image.open(image_path).convert(\"RGB\"), np.float32) / 255.0\n", 359 | "imshow(input_image, \"Input Image\")\n", 360 | "\n", 361 | "plt.subplot(1, 2, 2)\n", 362 | "imshow(final_pred_image, \"Predicted Image\")" 363 | ] 364 | } 365 | ], 366 | "metadata": { 367 | "accelerator": "GPU", 368 | "colab": { 369 | "provenance": [], 370 | "include_colab_link": true 371 | }, 372 | "kernelspec": { 373 | "display_name": "Python 3 (ipykernel)", 374 | "language": "python", 375 | "name": "python3" 376 | }, 377 | "language_info": { 378 | "codemirror_mode": { 379 | "name": "ipython", 380 | "version": 3 381 | }, 382 | "file_extension": ".py", 383 | "mimetype": "text/x-python", 384 | "name": "python", 385 | "nbconvert_exporter": "python", 386 | "pygments_lexer": "ipython3", 387 | "version": "3.8.2" 388 | } 389 | }, 390 | "nbformat": 4, 391 | "nbformat_minor": 0 392 | } -------------------------------------------------------------------------------- /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 | """Modified from https://github.com/google-research/maxim/blob/main/maxim/run_eval.py""" 16 | 17 | import os 18 | 19 | import numpy as np 20 | import tensorflow as tf 21 | from absl import app, flags 22 | from PIL import Image 23 | 24 | from create_maxim_model import Model 25 | from maxim.configs import MAXIM_CONFIGS 26 | 27 | FLAGS = flags.FLAGS 28 | 29 | flags.DEFINE_enum( 30 | "task", 31 | "Denoising", 32 | ["Denoising", "Deblurring", "Deraining", "Dehazing", "Enhancement"], 33 | "Task to run.", 34 | ) 35 | flags.DEFINE_string("ckpt_path", "", "Path to checkpoint.") 36 | flags.DEFINE_string("input_dir", "", "Input dir to the test set.") 37 | flags.DEFINE_string("output_dir", "", "Output dir to store predicted images.") 38 | flags.DEFINE_boolean("has_target", True, "Whether has corresponding gt image.") 39 | flags.DEFINE_boolean("save_images", True, "Dump predicted images.") 40 | flags.DEFINE_boolean("geometric_ensemble", False, "Whether use ensemble infernce.") 41 | 42 | _MODEL_VARIANT_DICT = { 43 | "Denoising": "S-3", 44 | "Deblurring": "S-3", 45 | "Deraining": "S-2", 46 | "Dehazing": "S-2", 47 | "Enhancement": "S-2", 48 | } 49 | 50 | 51 | _IMG_SIZE = 256 52 | 53 | _VALID_IMG_EXT = ["jpeg", "jpg", "png", "gif"] 54 | 55 | 56 | def mod_padding_symmetric(image, factor=64): 57 | """Padding the image to be divided by factor.""" 58 | height, width = image.shape[0], image.shape[1] 59 | height_pad, width_pad = ((height + factor) // factor) * factor, ( 60 | (width + factor) // factor 61 | ) * factor 62 | padh = height_pad - height if height % factor != 0 else 0 63 | padw = width_pad - width if width % factor != 0 else 0 64 | image = tf.pad( 65 | image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)], mode="REFLECT" 66 | ) 67 | return image 68 | 69 | 70 | # Since the model was not initialized to take variable-length sizes (None, None, 3), 71 | # we need to be careful about how we are resizing the images. 72 | # From https://www.tensorflow.org/lite/examples/style_transfer/overview#pre-process_the_inputs 73 | def resize_image(image, target_dim): 74 | # Resize the image so that the shorter dimension becomes `target_dim`. 75 | shape = tf.cast(tf.shape(image)[1:-1], tf.float32) 76 | short_dim = min(shape) 77 | scale = target_dim / short_dim 78 | new_shape = tf.cast(shape * scale, tf.int32) 79 | image = tf.image.resize(image, new_shape) 80 | 81 | # Central crop the image. 82 | image = tf.image.resize_with_crop_or_pad(image, target_dim, target_dim) 83 | 84 | return image 85 | 86 | 87 | def calculate_psnr(img1, img2, crop_border, test_y_channel=False): 88 | """Calculate PSNR (Peak Signal-to-Noise Ratio). 89 | 90 | Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio 91 | Args: 92 | img1 (ndarray): Images with range [0, 255]. 93 | img2 (ndarray): Images with range [0, 255]. 94 | crop_border (int): Cropped pixels in each edge of an image. These 95 | pixels are not involved in the PSNR calculation. 96 | test_y_channel (bool): Test on Y channel of YCbCr. Default: False. 97 | Returns: 98 | float: psnr result. 99 | """ 100 | assert ( 101 | img1.shape == img2.shape 102 | ), f"Image shapes are differnet: {img1.shape}, {img2.shape}." 103 | img1 = img1.astype(np.float64) 104 | img2 = img2.astype(np.float64) 105 | 106 | if crop_border != 0: 107 | img1 = img1[crop_border:-crop_border, crop_border:-crop_border, ...] 108 | img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...] 109 | 110 | if test_y_channel: 111 | img1 = to_y_channel(img1) 112 | img2 = to_y_channel(img2) 113 | 114 | mse = np.mean((img1 - img2) ** 2) 115 | if mse == 0: 116 | return float("inf") 117 | return 20.0 * np.log10(255.0 / np.sqrt(mse)) 118 | 119 | 120 | def _convert_input_type_range(img): 121 | """Convert the type and range of the input image. 122 | 123 | It converts the input image to np.float32 type and range of [0, 1]. 124 | It is mainly used for pre-processing the input image in colorspace 125 | convertion functions such as rgb2ycbcr and ycbcr2rgb. 126 | Args: 127 | img (ndarray): The input image. It accepts: 128 | 1. np.uint8 type with range [0, 255]; 129 | 2. np.float32 type with range [0, 1]. 130 | Returns: 131 | (ndarray): The converted image with type of np.float32 and range of 132 | [0, 1]. 133 | """ 134 | img_type = img.dtype 135 | img = img.astype(np.float32) 136 | if img_type == np.float32: 137 | pass 138 | elif img_type == np.uint8: 139 | img /= 255.0 140 | else: 141 | raise TypeError( 142 | "The img type should be np.float32 or np.uint8, " f"but got {img_type}" 143 | ) 144 | return img 145 | 146 | 147 | def _convert_output_type_range(img, dst_type): 148 | """Convert the type and range of the image according to dst_type. 149 | 150 | It converts the image to desired type and range. If `dst_type` is np.uint8, 151 | images will be converted to np.uint8 type with range [0, 255]. If 152 | `dst_type` is np.float32, it converts the image to np.float32 type with 153 | range [0, 1]. 154 | It is mainly used for post-processing images in colorspace convertion 155 | functions such as rgb2ycbcr and ycbcr2rgb. 156 | Args: 157 | img (ndarray): The image to be converted with np.float32 type and 158 | range [0, 255]. 159 | dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it 160 | converts the image to np.uint8 type with range [0, 255]. If 161 | dst_type is np.float32, it converts the image to np.float32 type 162 | with range [0, 1]. 163 | Returns: 164 | (ndarray): The converted image with desired type and range. 165 | """ 166 | if dst_type not in (np.uint8, np.float32): 167 | raise TypeError( 168 | "The dst_type should be np.float32 or np.uint8, " f"but got {dst_type}" 169 | ) 170 | if dst_type == np.uint8: 171 | img = img.round() 172 | else: 173 | img /= 255.0 174 | 175 | return img.astype(dst_type) 176 | 177 | 178 | def rgb2ycbcr(img, y_only=False): 179 | """Convert a RGB image to YCbCr image. 180 | 181 | This function produces the same results as Matlab's `rgb2ycbcr` function. 182 | It implements the ITU-R BT.601 conversion for standard-definition 183 | television. See more details in 184 | https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion. 185 | It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`. 186 | In OpenCV, it implements a JPEG conversion. See more details in 187 | https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion. 188 | 189 | Args: 190 | img (ndarray): The input image. It accepts: 191 | 1. np.uint8 type with range [0, 255]; 192 | 2. np.float32 type with range [0, 1]. 193 | y_only (bool): Whether to only return Y channel. Default: False. 194 | Returns: 195 | ndarray: The converted YCbCr image. The output image has the same type 196 | and range as input image. 197 | """ 198 | img_type = img.dtype 199 | img = _convert_input_type_range(img) 200 | if y_only: 201 | out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0 202 | else: 203 | out_img = np.matmul( 204 | img, 205 | [ 206 | [65.481, -37.797, 112.0], 207 | [128.553, -74.203, -93.786], 208 | [24.966, 112.0, -18.214], 209 | ], 210 | ) + [16, 128, 128] 211 | out_img = _convert_output_type_range(out_img, img_type) 212 | return out_img 213 | 214 | 215 | def to_y_channel(img): 216 | """Change to Y channel of YCbCr. 217 | 218 | Args: 219 | img (ndarray): Images with range [0, 255]. 220 | Returns: 221 | (ndarray): Images with range [0, 255] (float type) without round. 222 | """ 223 | img = img.astype(np.float32) / 255.0 224 | if img.ndim == 3 and img.shape[2] == 3: 225 | img = rgb2ycbcr(img, y_only=True) 226 | img = img[..., None] 227 | return img * 255.0 228 | 229 | 230 | def augment_image(image, times=8): 231 | """Geometric augmentation.""" 232 | if times == 4: # only rotate image 233 | images = [] 234 | for k in range(0, 4): 235 | images.append(np.rot90(image, k=k)) 236 | images = np.stack(images, axis=0) 237 | elif times == 8: # roate and flip image 238 | images = [] 239 | for k in range(0, 4): 240 | images.append(np.rot90(image, k=k)) 241 | image = np.fliplr(image) 242 | for k in range(0, 4): 243 | images.append(np.rot90(image, k=k)) 244 | images = np.stack(images, axis=0) 245 | else: 246 | raise Exception(f"Error times: {times}") 247 | return images 248 | 249 | 250 | def deaugment_image(images, times=8): 251 | """Reverse the geometric augmentation.""" 252 | 253 | if times == 4: # only rotate image 254 | image = [] 255 | for k in range(0, 4): 256 | image.append(np.rot90(images[k], k=4 - k)) 257 | image = np.stack(image, axis=0) 258 | image = np.mean(image, axis=0) 259 | elif times == 8: # roate and flip image 260 | image = [] 261 | for k in range(0, 4): 262 | image.append(np.rot90(images[k], k=4 - k)) 263 | for k in range(0, 4): 264 | image.append(np.fliplr(np.rot90(images[4 + k], k=4 - k))) 265 | image = np.mean(image, axis=0) 266 | else: 267 | raise Exception(f"Error times: {times}") 268 | return image 269 | 270 | 271 | def is_image_file(filename): 272 | """Check if it is an valid image file by extension.""" 273 | 274 | return any( 275 | (filename.endswith(extension)) or (filename.endswith(extension.upper())) 276 | for extension in _VALID_IMG_EXT 277 | ) 278 | 279 | 280 | def save_img(img, pth): 281 | """Save an image to disk. 282 | 283 | Args: 284 | img: np.ndarry, [height, width, channels], img will be clipped to [0, 1] 285 | before saved to pth. 286 | pth: string, path to save the image to. 287 | """ 288 | Image.fromarray(np.array((np.clip(img, 0.0, 1.0) * 255.0).astype(np.uint8))).save( 289 | pth, "PNG" 290 | ) 291 | 292 | 293 | def make_shape_even(image): 294 | """Pad the image to have even shapes.""" 295 | height, width = image.shape[0], image.shape[1] 296 | padh = 1 if height % 2 != 0 else 0 297 | padw = 1 if width % 2 != 0 else 0 298 | image = tf.pad(image, [(0, padh), (0, padw), (0, 0)], mode="REFLECT") 299 | return image 300 | 301 | 302 | def main(_): 303 | if FLAGS.save_images: 304 | os.makedirs(FLAGS.output_dir, exist_ok=True) 305 | 306 | # sorted is important for continuning an inference job. 307 | filepath = sorted(os.listdir(os.path.join(FLAGS.input_dir, "input"))) 308 | input_filenames = [ 309 | os.path.join(FLAGS.input_dir, "input", x) for x in filepath if is_image_file(x) 310 | ] 311 | if FLAGS.has_target: 312 | target_filenames = [ 313 | os.path.join(FLAGS.input_dir, "target", x) 314 | for x in filepath 315 | if is_image_file(x) 316 | ] 317 | num_images = len(input_filenames) 318 | 319 | print("Initializing model and loading model weights.") 320 | model = tf.keras.models.load_model(FLAGS.ckpt_path) 321 | print("Model successfully initialized and weights loaded.") 322 | 323 | psnr_all = [] 324 | 325 | def _process_file(i): 326 | print(f"Processing {i + 1} / {num_images}...") 327 | input_file = input_filenames[i] 328 | input_img = np.asarray(Image.open(input_file).convert("RGB"), np.float32) / 255.0 329 | 330 | if FLAGS.has_target: 331 | target_file = target_filenames[i] 332 | target_img = ( 333 | np.asarray(Image.open(target_file).convert("RGB"), np.float32) / 255.0 334 | ) 335 | 336 | height, width = input_img.shape[0], input_img.shape[1] 337 | # Padding images to have even shapes 338 | input_img = make_shape_even(input_img) 339 | height_even, width_even = input_img.shape[0], input_img.shape[1] 340 | 341 | # padding images to be multiplies of 64 342 | input_img = mod_padding_symmetric(input_img, factor=64) 343 | 344 | if FLAGS.geometric_ensemble: 345 | input_img = augment_image(input_img, FLAGS.ensemble_times) 346 | else: 347 | input_img = tf.expand_dims(input_img, axis=0) 348 | 349 | # handle multi-stage outputs, obtain the last scale output of last stage 350 | 351 | preds = model.predict(input_img) 352 | if isinstance(preds, list): 353 | preds = preds[-1] 354 | if isinstance(preds, list): 355 | preds = preds[-1] 356 | 357 | # De-ensemble by averaging inferenced results. 358 | if FLAGS.geometric_ensemble: 359 | preds = deaugment_image(preds, FLAGS.ensemble_times) 360 | else: 361 | preds = np.array(preds[0], np.float32) 362 | 363 | # unpad images to get the original resolution 364 | new_height, new_width = preds.shape[0], preds.shape[1] 365 | h_start = new_height // 2 - height_even // 2 366 | h_end = h_start + height 367 | w_start = new_width // 2 - width_even // 2 368 | w_end = w_start + width 369 | preds = preds[h_start:h_end, w_start:w_end, :] 370 | 371 | # print PSNR scores 372 | if FLAGS.has_target: 373 | psnr = calculate_psnr( 374 | target_img * 255.0, preds * 255.0, crop_border=0, test_y_channel=False 375 | ) 376 | print(f"{i}th image: psnr = {psnr:.4f}") 377 | else: 378 | psnr = -1 379 | 380 | # save files 381 | basename = os.path.basename(input_file) 382 | if FLAGS.save_images: 383 | save_pth = os.path.join(FLAGS.output_dir, basename) 384 | save_img(preds, save_pth) 385 | 386 | return psnr 387 | 388 | for i in range(num_images): 389 | psnr = _process_file(i) 390 | psnr_all.append(psnr) 391 | 392 | psnr_all = np.asarray(psnr_all) 393 | 394 | print(f"average psnr = {np.sum(psnr_all)/num_images:.4f}") 395 | print(f"std psnr = {np.std(psnr_all):.4f}") 396 | 397 | 398 | if __name__ == "__main__": 399 | app.run(main) 400 | -------------------------------------------------------------------------------- /maxim/maxim.py: -------------------------------------------------------------------------------- 1 | """ 2 | MAXIM based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py 3 | """ 4 | 5 | import functools 6 | 7 | import tensorflow as tf 8 | from tensorflow.keras import backend as K 9 | from tensorflow.keras import layers 10 | 11 | from .blocks.attentions import SAM 12 | from .blocks.bottleneck import BottleneckBlock 13 | from .blocks.misc_gating import CrossGatingBlock 14 | from .blocks.others import UpSampleRatio 15 | from .blocks.unet import UNetDecoderBlock, UNetEncoderBlock 16 | from .layers import Resizing 17 | 18 | Conv1x1 = functools.partial(layers.Conv2D, kernel_size=(1, 1), padding="same") 19 | Conv3x3 = functools.partial(layers.Conv2D, kernel_size=(3, 3), padding="same") 20 | ConvT_up = functools.partial( 21 | layers.Conv2DTranspose, kernel_size=(2, 2), strides=(2, 2), padding="same" 22 | ) 23 | Conv_down = functools.partial( 24 | layers.Conv2D, kernel_size=(4, 4), strides=(2, 2), padding="same" 25 | ) 26 | 27 | 28 | def MAXIM( 29 | features: int = 64, 30 | depth: int = 3, 31 | num_stages: int = 2, 32 | num_groups: int = 1, 33 | use_bias: bool = True, 34 | num_supervision_scales: int = 1, 35 | lrelu_slope: float = 0.2, 36 | use_global_mlp: bool = True, 37 | use_cross_gating: bool = True, 38 | high_res_stages: int = 2, 39 | block_size_hr=(16, 16), 40 | block_size_lr=(8, 8), 41 | grid_size_hr=(16, 16), 42 | grid_size_lr=(8, 8), 43 | num_bottleneck_blocks: int = 1, 44 | block_gmlp_factor: int = 2, 45 | grid_gmlp_factor: int = 2, 46 | input_proj_factor: int = 2, 47 | channels_reduction: int = 4, 48 | num_outputs: int = 3, 49 | dropout_rate: float = 0.0, 50 | ): 51 | """The MAXIM model function with multi-stage and multi-scale supervision. 52 | 53 | For more model details, please check the CVPR paper: 54 | MAXIM: MUlti-Axis MLP for Image Processing (https://arxiv.org/abs/2201.02973) 55 | 56 | Attributes: 57 | features: initial hidden dimension for the input resolution. 58 | depth: the number of downsampling depth for the model. 59 | num_stages: how many stages to use. It will also affects the output list. 60 | num_groups: how many blocks each stage contains. 61 | use_bias: whether to use bias in all the conv/mlp layers. 62 | num_supervision_scales: the number of desired supervision scales. 63 | lrelu_slope: the negative slope parameter in leaky_relu layers. 64 | use_global_mlp: whether to use the multi-axis gated MLP block (MAB) in each 65 | layer. 66 | use_cross_gating: whether to use the cross-gating MLP block (CGB) in the 67 | skip connections and multi-stage feature fusion layers. 68 | high_res_stages: how many stages are specificied as high-res stages. The 69 | rest (depth - high_res_stages) are called low_res_stages. 70 | block_size_hr: the block_size parameter for high-res stages. 71 | block_size_lr: the block_size parameter for low-res stages. 72 | grid_size_hr: the grid_size parameter for high-res stages. 73 | grid_size_lr: the grid_size parameter for low-res stages. 74 | num_bottleneck_blocks: how many bottleneck blocks. 75 | block_gmlp_factor: the input projection factor for block_gMLP layers. 76 | grid_gmlp_factor: the input projection factor for grid_gMLP layers. 77 | input_proj_factor: the input projection factor for the MAB block. 78 | channels_reduction: the channel reduction factor for SE layer. 79 | num_outputs: the output channels. 80 | dropout_rate: Dropout rate. 81 | 82 | Returns: 83 | The output contains a list of arrays consisting of multi-stage multi-scale 84 | outputs. For example, if num_stages = num_supervision_scales = 3 (the 85 | model used in the paper), the output specs are: outputs = 86 | [[output_stage1_scale1, output_stage1_scale2, output_stage1_scale3], 87 | [output_stage2_scale1, output_stage2_scale2, output_stage2_scale3], 88 | [output_stage3_scale1, output_stage3_scale2, output_stage3_scale3],] 89 | The final output can be retrieved by outputs[-1][-1]. 90 | """ 91 | 92 | def apply(x): 93 | n, h, w, c = ( 94 | K.int_shape(x)[0], 95 | K.int_shape(x)[1], 96 | K.int_shape(x)[2], 97 | K.int_shape(x)[3], 98 | ) # input image shape 99 | 100 | shortcuts = [] 101 | shortcuts.append(x) 102 | 103 | # Get multi-scale input images 104 | for i in range(1, num_supervision_scales): 105 | resizing_layer = Resizing( 106 | ratio=(2**i), 107 | method="nearest", 108 | antialias=True, # Following `jax.image.resize()`. 109 | name=f"initial_resizing_{K.get_uid('Resizing')}", 110 | ) 111 | shortcuts.append(resizing_layer(x)) 112 | 113 | # store outputs from all stages and all scales 114 | # Eg, [[(64, 64, 3), (128, 128, 3), (256, 256, 3)], # Stage-1 outputs 115 | # [(64, 64, 3), (128, 128, 3), (256, 256, 3)],] # Stage-2 outputs 116 | outputs_all = [] 117 | sam_features, encs_prev, decs_prev = [], [], [] 118 | 119 | for idx_stage in range(num_stages): 120 | # Input convolution, get multi-scale input features 121 | x_scales = [] 122 | for i in range(num_supervision_scales): 123 | x_scale = Conv3x3( 124 | filters=(2**i) * features, 125 | use_bias=use_bias, 126 | name=f"stage_{idx_stage}_input_conv_{i}", 127 | )(shortcuts[i]) 128 | 129 | # If later stages, fuse input features with SAM features from prev stage 130 | if idx_stage > 0: 131 | # use larger blocksize at high-res stages 132 | if use_cross_gating: 133 | block_size = ( 134 | block_size_hr if i < high_res_stages else block_size_lr 135 | ) 136 | grid_size = ( 137 | grid_size_hr if i < high_res_stages else block_size_lr 138 | ) 139 | x_scale, _ = CrossGatingBlock( 140 | features=(2**i) * features, 141 | block_size=block_size, 142 | grid_size=grid_size, 143 | dropout_rate=dropout_rate, 144 | input_proj_factor=input_proj_factor, 145 | upsample_y=False, 146 | use_bias=use_bias, 147 | name=f"stage_{idx_stage}_input_fuse_sam_{i}", 148 | )(x_scale, sam_features.pop()) 149 | else: 150 | x_scale = Conv1x1( 151 | filters=(2**i) * features, 152 | use_bias=use_bias, 153 | name=f"stage_{idx_stage}_input_catconv_{i}", 154 | )(tf.concat([x_scale, sam_features.pop()], axis=-1)) 155 | 156 | x_scales.append(x_scale) 157 | 158 | # start encoder blocks 159 | encs = [] 160 | x = x_scales[0] # First full-scale input feature 161 | 162 | for i in range(depth): # 0, 1, 2 163 | # use larger blocksize at high-res stages, vice versa. 164 | block_size = block_size_hr if i < high_res_stages else block_size_lr 165 | grid_size = grid_size_hr if i < high_res_stages else block_size_lr 166 | use_cross_gating_layer = True if idx_stage > 0 else False 167 | 168 | # Multi-scale input if multi-scale supervision 169 | x_scale = x_scales[i] if i < num_supervision_scales else None 170 | 171 | # UNet Encoder block 172 | enc_prev = encs_prev.pop() if idx_stage > 0 else None 173 | dec_prev = decs_prev.pop() if idx_stage > 0 else None 174 | 175 | x, bridge = UNetEncoderBlock( 176 | num_channels=(2**i) * features, 177 | num_groups=num_groups, 178 | downsample=True, 179 | lrelu_slope=lrelu_slope, 180 | block_size=block_size, 181 | grid_size=grid_size, 182 | block_gmlp_factor=block_gmlp_factor, 183 | grid_gmlp_factor=grid_gmlp_factor, 184 | input_proj_factor=input_proj_factor, 185 | channels_reduction=channels_reduction, 186 | use_global_mlp=use_global_mlp, 187 | dropout_rate=dropout_rate, 188 | use_bias=use_bias, 189 | use_cross_gating=use_cross_gating_layer, 190 | name=f"stage_{idx_stage}_encoder_block_{i}", 191 | )(x, skip=x_scale, enc=enc_prev, dec=dec_prev) 192 | 193 | # Cache skip signals 194 | encs.append(bridge) 195 | 196 | # Global MLP bottleneck blocks 197 | for i in range(num_bottleneck_blocks): 198 | x = BottleneckBlock( 199 | block_size=block_size_lr, 200 | grid_size=block_size_lr, 201 | features=(2 ** (depth - 1)) * features, 202 | num_groups=num_groups, 203 | block_gmlp_factor=block_gmlp_factor, 204 | grid_gmlp_factor=grid_gmlp_factor, 205 | input_proj_factor=input_proj_factor, 206 | dropout_rate=dropout_rate, 207 | use_bias=use_bias, 208 | channels_reduction=channels_reduction, 209 | name=f"stage_{idx_stage}_global_block_{i}", 210 | )(x) 211 | # cache global feature for cross-gating 212 | global_feature = x 213 | 214 | # start cross gating. Use multi-scale feature fusion 215 | skip_features = [] 216 | for i in reversed(range(depth)): # 2, 1, 0 217 | # use larger blocksize at high-res stages 218 | block_size = block_size_hr if i < high_res_stages else block_size_lr 219 | grid_size = grid_size_hr if i < high_res_stages else block_size_lr 220 | 221 | # get additional multi-scale signals 222 | signal = tf.concat( 223 | [ 224 | UpSampleRatio( 225 | num_channels=(2**i) * features, 226 | ratio=2 ** (j - i), 227 | use_bias=use_bias, 228 | name=f"UpSampleRatio_{K.get_uid('UpSampleRatio')}", 229 | )(enc) 230 | for j, enc in enumerate(encs) 231 | ], 232 | axis=-1, 233 | ) 234 | 235 | # Use cross-gating to cross modulate features 236 | if use_cross_gating: 237 | skips, global_feature = CrossGatingBlock( 238 | features=(2**i) * features, 239 | block_size=block_size, 240 | grid_size=grid_size, 241 | input_proj_factor=input_proj_factor, 242 | dropout_rate=dropout_rate, 243 | upsample_y=True, 244 | use_bias=use_bias, 245 | name=f"stage_{idx_stage}_cross_gating_block_{i}", 246 | )(signal, global_feature) 247 | else: 248 | skips = Conv1x1( 249 | filters=(2**i) * features, use_bias=use_bias, name="Conv_0" 250 | )(signal) 251 | skips = Conv3x3( 252 | filters=(2**i) * features, use_bias=use_bias, name="Conv_1" 253 | )(skips) 254 | 255 | skip_features.append(skips) 256 | 257 | # start decoder. Multi-scale feature fusion of cross-gated features 258 | outputs, decs, sam_features = [], [], [] 259 | for i in reversed(range(depth)): 260 | # use larger blocksize at high-res stages 261 | block_size = block_size_hr if i < high_res_stages else block_size_lr 262 | grid_size = grid_size_hr if i < high_res_stages else block_size_lr 263 | 264 | # get multi-scale skip signals from cross-gating block 265 | signal = tf.concat( 266 | [ 267 | UpSampleRatio( 268 | num_channels=(2**i) * features, 269 | ratio=2 ** (depth - j - 1 - i), 270 | use_bias=use_bias, 271 | name=f"UpSampleRatio_{K.get_uid('UpSampleRatio')}", 272 | )(skip) 273 | for j, skip in enumerate(skip_features) 274 | ], 275 | axis=-1, 276 | ) 277 | 278 | # Decoder block 279 | x = UNetDecoderBlock( 280 | num_channels=(2**i) * features, 281 | num_groups=num_groups, 282 | lrelu_slope=lrelu_slope, 283 | block_size=block_size, 284 | grid_size=grid_size, 285 | block_gmlp_factor=block_gmlp_factor, 286 | grid_gmlp_factor=grid_gmlp_factor, 287 | input_proj_factor=input_proj_factor, 288 | channels_reduction=channels_reduction, 289 | use_global_mlp=use_global_mlp, 290 | dropout_rate=dropout_rate, 291 | use_bias=use_bias, 292 | name=f"stage_{idx_stage}_decoder_block_{i}", 293 | )(x, bridge=signal) 294 | 295 | # Cache decoder features for later-stage's usage 296 | decs.append(x) 297 | 298 | # output conv, if not final stage, use supervised-attention-block. 299 | if i < num_supervision_scales: 300 | if idx_stage < num_stages - 1: # not last stage, apply SAM 301 | sam, output = SAM( 302 | num_channels=(2**i) * features, 303 | output_channels=num_outputs, 304 | use_bias=use_bias, 305 | name=f"stage_{idx_stage}_supervised_attention_module_{i}", 306 | )(x, shortcuts[i]) 307 | outputs.append(output) 308 | sam_features.append(sam) 309 | else: # Last stage, apply output convolutions 310 | output = Conv3x3( 311 | num_outputs, 312 | use_bias=use_bias, 313 | name=f"stage_{idx_stage}_output_conv_{i}", 314 | )(x) 315 | output = output + shortcuts[i] 316 | outputs.append(output) 317 | # Cache encoder and decoder features for later-stage's usage 318 | encs_prev = encs[::-1] 319 | decs_prev = decs 320 | 321 | # Store outputs 322 | outputs_all.append(outputs) 323 | return outputs_all 324 | 325 | return apply 326 | --------------------------------------------------------------------------------