Update app.py

app.py

@@ -1,82 +1,41 @@
 '''NEURAL STYLE TRANSFER '''
 
-import gradio as gr
+import numpy as np
 import tensorflow as tf
 import tensorflow_hub as hub
-import PIL
+import gradio as gr
 from PIL import Image
-import numpy as np
-# import time
-# import requests   
-#import cv2
-
-# !mkdir nstmodel 
-# !wget -c https://storage.googleapis.com/tfhub-modules/google/magenta/arbitrary-image-stylization-v1-256/2.tar.gz -O - | tar -xz -C /nstmodel
-# import tensorflow.keras
 
-# from PIL import Image, ImageOps
-#import requests
-#import tarfile
-
-#MODEL_PATH='Nst_model'
-
-# Disable scientific notation for clarity
 np.set_printoptions(suppress=True)
 
-# Load model from TF-Hub
-model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2') 
-
-# Load the model
-#model = tf.keras.models.load_model(MODEL_PATH)
+model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
 
 def tensor_to_image(tensor):
-  tensor = tensor*255
-  tensor = np.array(tensor, dtype=np.uint8)
-  if np.ndim(tensor)>3:
-    assert tensor.shape[0] == 1
-    tensor = tensor[0]
-  return PIL.Image.fromarray(tensor)
-
-
-
-"""## Grayscaling image for testing purpose to check if we could get better results.
-def gray_scaled(inp_img):
-  gray = cv2.cvtColor(inp_img, cv2.COLOR_BGR2GRAY)
-  gray_img = np.zeros_like(inp_img)
-  gray_img[:,:,0] = gray
-  gray_img[:,:,1] = gray
-  gray_img[:,:,2] = gray
-  return gray_img
- """
-
-##Transformation 
-def transform_my_model(content_image,style_image):
-  # Convert to float32 numpy array, add batch dimension, and normalize to range [0, 1]
-  #content_image=gray_scaled(content_image)
-  content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.
-  style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.
- 
-  #Resizing image
-  #style_image = tf.image.resize(style_image, (256, 256))
-
-  # Stylize image
-  outputs = model(tf.constant(content_image), tf.constant(style_image))
-  stylized_image = outputs[0]
-
-  # stylized = tf.image.resize(stylized_image, (356, 356))
-  stylized_image =tensor_to_image(stylized_image)
-  return stylized_image
-
-
-image1 = gr.Image(label="Content Image") #CONTENT IMAGE
-image2 = gr.Image(label="Style Image") #STYLE IMAGE
-stylizedimg=gr.Image(label="Result") 
-gr.Interface(fn=transform_my_model, inputs= [image1,image2] , 
-             outputs= stylizedimg,title='Style Transfer',theme='seafoam',
-             examples=[['Content_Images/contnt12.jpg','VG516.jpg'],['Content_Images/contnt2.jpg','Content_Images/styl9.jpg'],['Content_Images/contnt.jpg','Content_Images/styl22.jpg']],
-             article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin.",
-             share=True
-            ).launch()
-
-
-
+    tensor *= 255
+    tensor = np.array(tensor, dtype=np.uint8)
+    if tensor.ndim > 3:
+        tensor = tensor[0]
+    return Image.fromarray(tensor)
+
+def transform_my_model(content_image, style_image):
+    content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
+    style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
+    stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
+    return tensor_to_image(stylized_image)
+
+demo = gr.Interface(
+    fn=transform_my_model,
+    inputs=[gr.Image(label="Content Image"), gr.Image(label="Style Image")],
+    outputs=[gr.Image(label="Result")],
+    title="Style Transfer",
+    theme="seafoam",
+    examples=[
+        ["Content_Images/contnt12.jpg", "VG516.jpg"],
+        ["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
+        ["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
+    ],
+    article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin.",
+    share=True
+)
+
+demo.launch()