Update app.py

app.py

@@ -5,34 +5,14 @@ import numpy as np
 import cv2
 import gradio as gr
 import os
-import subprocess
-import gdown
-
-os.makedirs('/content', exist_ok=True)
-# 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_path = './model/wt_generator_best.keras'
 Pix2Pix = keras.saving.load_model(pix2pix_path)
 
 # Load MS-UNet
-unet_path = '/model/unet.keras'
+unet_path = './model/unet.keras'
 MS_UNet = keras.saving.load_model(unet_path)
 
 # Load WD-Net
@@ -43,16 +23,66 @@ class Clip(keras.layers.Layer):
   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)
+old_gen_path = './model/generator_epoch_10.keras'
+WD_Net_old = keras.saving.load_model(old_gen_path)
+new_gen_path = './model/WD-Net_generator.keras'
+WD_Net_new = keras.saving.load_model(new_gen_path)
 
 # Define infer function
 def infer(img, model='WD-Net'):
+  # Read image
+  # img = tf.image.decode_png(tf.io.read_file('./data/' + img_path), channels=3)
+  # Image original shape
+  org_shape = img.shape
+  org_img = tf.image.resize(img, [256, 256], method='area')
+  org_img = tf.cast(org_img, tf.uint8).numpy()
+  org_img = cv2.resize(org_img, (org_shape[1], org_shape[0]))
+
+  # Choose model
+  if model == 'WD-Net':
+    generator = WD_Net_old
+    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 org_img, out_img
+
+def infer_v1(img_path, model="WD_Net"):
+  # Read image
+  img = tf.image.decode_png(tf.io.read_file('./data/' + img_path), channels=3)
   # Image original shape
   org_shape = img.shape
+  # org_img = tf.image.resize(img, [256, 256], method='area')
+  org_img = tf.cast(img, tf.uint8).numpy()
+  # org_img = cv2.resize(org_img, (org_shape[1], org_shape[0]))
+
   # Choose model
   if model == 'WD-Net':
-    generator = WD_Net
+    generator = WD_Net_new
     img = tf.image.resize(img, [256, 256], method='area')
     # Normalize image and return
     img = tf.cast(img, tf.float32) / 255.
@@ -81,15 +111,28 @@ def infer(img, model='WD-Net'):
     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
+  return org_img, out_img
 
 # Main gradio code
+# Define data and sort it
+data = os.listdir('./data')
+data.sort()
+
+# Model list
 model_list = ['WD-Net', 'MS-UNet', 'Pix2Pix']
 
 demo = gr.Interface(
     fn=infer,
-    inputs=[gr.Image(), gr.Dropdown(model_list)],
-    outputs=gr.Image(),
+    inputs=[gr.Image(label="Choose an Image"), gr.Dropdown(model_list, label="Model")],
+    outputs=[gr.Image(label="Watermarked Image"), gr.Image(label="Removed Watermarked Image")],
 )
 
-demo.launch()
+demo_v1 = gr.Interface(
+    fn=infer_v1,
+    inputs=[gr.Dropdown(data, label="Choose an Image"), gr.Dropdown(model_list, label="Model")],
+    outputs=[gr.Image(label="Watermarked Image"), gr.Image(label="Removed Watermarked Image")],
+)
+
+tabbed_interface = gr.TabbedInterface([demo, demo_v1], ["Document", "Patch"], title="Watermark Removal")
+
+tabbed_interface.launch()