Create new file

app.py

@@ -0,0 +1,144 @@
+from huggingface_hub.keras_mixin import from_pretrained_keras
+
+from PIL import Image
+
+import numpy as np
+
+from create_maxim_model import Model
+from maxim.configs import MAXIM_CONFIGS
+
+
+_MODEL = from_pretrained_keras("sayakpaul/S-2_enhancement_lol")
+
+
+def mod_padding_symmetric(image, factor=64):
+    """Padding the image to be divided by factor."""
+    height, width = image.shape[0], image.shape[1]
+    height_pad, width_pad = ((height + factor) // factor) * factor, (
+        (width + factor) // factor
+    ) * factor
+    padh = height_pad - height if height % factor != 0 else 0
+    padw = width_pad - width if width % factor != 0 else 0
+    image = tf.pad(
+        image, [(padh // 2, padh // 2), (padw // 2, padw // 2), (0, 0)], mode="REFLECT"
+    )
+    return image
+    
+ def _convert_input_type_range(img):
+    """Convert the type and range of the input image.
+
+    It converts the input image to np.float32 type and range of [0, 1].
+    It is mainly used for pre-processing the input image in colorspace
+    convertion functions such as rgb2ycbcr and ycbcr2rgb.
+    Args:
+      img (ndarray): The input image. It accepts:
+          1. np.uint8 type with range [0, 255];
+          2. np.float32 type with range [0, 1].
+    Returns:
+        (ndarray): The converted image with type of np.float32 and range of
+            [0, 1].
+    """
+    img_type = img.dtype
+    img = img.astype(np.float32)
+    if img_type == np.float32:
+        pass
+    elif img_type == np.uint8:
+        img /= 255.0
+    else:
+        raise TypeError(
+            "The img type should be np.float32 or np.uint8, " f"but got {img_type}"
+        )
+    return img
+
+
+def _convert_output_type_range(img, dst_type):
+    """Convert the type and range of the image according to dst_type.
+
+    It converts the image to desired type and range. If `dst_type` is np.uint8,
+    images will be converted to np.uint8 type with range [0, 255]. If
+    `dst_type` is np.float32, it converts the image to np.float32 type with
+    range [0, 1].
+    It is mainly used for post-processing images in colorspace convertion
+    functions such as rgb2ycbcr and ycbcr2rgb.
+    Args:
+      img (ndarray): The image to be converted with np.float32 type and
+          range [0, 255].
+      dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
+          converts the image to np.uint8 type with range [0, 255]. If
+          dst_type is np.float32, it converts the image to np.float32 type
+          with range [0, 1].
+    Returns:
+      (ndarray): The converted image with desired type and range.
+    """
+    if dst_type not in (np.uint8, np.float32):
+        raise TypeError(
+            "The dst_type should be np.float32 or np.uint8, " f"but got {dst_type}"
+        )
+    if dst_type == np.uint8:
+        img = img.round()
+    else:
+        img /= 255.0
+
+    return img.astype(dst_type)
+
+
+def make_shape_even(image):
+    """Pad the image to have even shapes."""
+    height, width = image.shape[0], image.shape[1]
+    padh = 1 if height % 2 != 0 else 0
+    padw = 1 if width % 2 != 0 else 0
+    image = tf.pad(image, [(0, padh), (0, padw), (0, 0)], mode="REFLECT")
+    return image
+    
+    
+def process_image(image: Image):
+    input_img = np.asarray(image) / 255.0
+    height, width = input_img.shape[0], input_img.shape[1]
+
+    # Padding images to have even shapes
+    input_img = make_shape_even(input_img)
+    height_even, width_even = input_img.shape[0], input_img.shape[1]
+
+    # padding images to be multiplies of 64
+    input_img = mod_padding_symmetric(input_img, factor=64)
+    input_img = tf.expand_dims(input_img, axis=0)
+    return input_img, height_even, width_even
+    
+   
+def init_new_model(input_img):
+    configs = MAXIM_CONFIGS.get("S-2")
+    configs.update(
+        {
+            "variant": "S-2",
+            "dropout_rate": 0.0,
+            "num_outputs": 3,
+            "use_bias": True,
+            "num_supervision_scales": 3,
+        }
+    )
+    configs.update({"input_resolution": (input_img.shape[1], input_img.shape[2])})
+    new_model = Model(**configs)
+    new_model.set_weights(_MODEL.get_weights())
+    return new_model
+    
+    
+def infer(image):
+    preprocessed_image, height_even, width_even = process_image(image)
+    new_model = init_new_model(preprocessed_image)
+    
+    preds = new_model.predict(preprocessed_image)
+    if isinstance(preds, list):
+        preds = preds[-1]
+        if isinstance(preds, list):
+            preds = preds[-1]
+    
+    preds = np.array(preds[0], np.float32)
+    
+    new_height, new_width = preds.shape[0], preds.shape[1]
+    h_start = new_height // 2 - height_even // 2
+    h_end = h_start + height
+    w_start = new_width // 2 - width_even // 2
+    w_end = w_start + width
+    preds = preds[h_start:h_end, w_start:w_end, :]
+    
+    return Image.fromarray(np.array((np.clip(preds, 0.0, 1.0) * 255.0).astype(np.uint8)))