Initial commit

app.py

@@ -0,0 +1,94 @@
+# Import all necessary libs
+import tensorflow as tf
+import keras
+import numpy as np
+import cv2
+import gradio as gr
+import os
+import subprocess
+import gdown
+
+# Download model
+# Pix2pix model
+model_url = 'https://drive.google.com/drive/folders/1jOxiKyf8n7fwNZfgeZyUrNJ90LEmIL3S?usp=sharing'
+os.makedirs('/content/pix2pix', exist_ok=True)
+subprocess.run(['gdown', '--fuzzy', model_url, '-O', '/content/pix2pix', '--folder'], check=True)
+
+# WD-Net
+model_url = 'https://drive.google.com/file/d/1M8EOE4Ej8oS4_0BHCEwExxu5CMFS5HQZ/view?usp=sharing'
+os.makedirs('/content/WD-Net', exist_ok=True)
+subprocess.run(['gdown', '--fuzzy', model_url, '-O', '/content/WD-Net/model.zip'], check=True)
+subprocess.run(['unzip', '/content/WD-Net/model.zip', '-d', '/content/WD-Net'], check=True)
+
+# MS-UNet
+model_url = 'https://drive.google.com/file/d/1-0_bEWTItkILbCJQ4ViEBGg0zJPaIcC1/view?usp=sharing'
+os.makedirs('/content/MS-UNet', exist_ok=True)
+subprocess.run(['gdown', '--fuzzy', model_url, '-O', '/content/MS-UNet/unet.keras'], check=True)
+
+# Load model
+# Load Pix2Pix
+pix2pix_path = '/model/wt_generator_best.keras'
+Pix2Pix = keras.saving.load_model(pix2pix_path)
+
+# Load MS-UNet
+unet_path = '/model/unet.keras'
+MS_UNet = keras.saving.load_model(unet_path)
+
+# Load WD-Net
+@keras.saving.register_keras_serializable(package="Clip")
+class Clip(keras.layers.Layer):
+  def __init__(self, **kwargs):
+    super().__init__(**kwargs)
+  def call(self, input):
+    return tf.clip_by_value(input, 0, 1)
+
+gen_path = '/model/generator_epoch:10.keras'
+WD_Net = keras.saving.load_model(gen_path)
+
+# Define infer function
+def infer(img, model='WD-Net'):
+  # Image original shape
+  org_shape = img.shape
+  # Choose model
+  if model == 'WD-Net':
+    generator = WD_Net
+    img = tf.image.resize(img, [256, 256], method='area')
+    # Normalize image and return
+    img = tf.cast(img, tf.float32) / 255.
+    img = tf.expand_dims(img, axis=0)
+    rm_wt = generator.predict(img, verbose=0)
+    rm_wt = rm_wt['I'][0]
+    rm_wt = cv2.resize(rm_wt, (org_shape[1], org_shape[0]))
+    out_img = (rm_wt * 255).astype(np.uint8)
+  elif model == 'MS-UNet':
+    generator = MS_UNet
+    img = tf.image.resize(img, [256, 256], method='area')
+    # Normalize image and return
+    img = (tf.cast(img, tf.float32) - 127.5) / 127.5
+    img = tf.expand_dims(img, axis=0)
+    rm_wt = generator.predict(img, verbose=0)
+    rm_wt = rm_wt[0]
+    rm_wt = cv2.resize(rm_wt, (org_shape[1], org_shape[0]))
+    out_img = ((rm_wt + 1) / 2 * 255).astype(np.uint8)
+  elif model == 'Pix2Pix':
+    generator = Pix2Pix
+    img = tf.image.resize(img, [256, 256], method='area')
+    # Normalize image and return
+    img = tf.cast(img, tf.float32) / 255.
+    img = tf.expand_dims(img, axis=0)
+    rm_wt = generator.predict(img, verbose=0)
+    rm_wt = rm_wt[0]
+    rm_wt = cv2.resize(rm_wt, (org_shape[1], org_shape[0]))
+    out_img = (rm_wt * 255).astype(np.uint8)
+  return out_img
+
+# Main gradio code
+model_list = ['WD-Net', 'MS-UNet', 'Pix2Pix']
+
+demo = gr.Interface(
+    fn=infer,
+    inputs=[gr.Image(), gr.Dropdown(model_list)],
+    outputs=gr.Image(),
+)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gdown
+tensorflow
+matplotlib
+opencv-python
+numpy
+pandas
+pillow