Update app.py

app.py

@@ -10,6 +10,7 @@ import tensorflow as tf
 import tensorflow_hub as hub
 from PIL import Image
 import numpy as np
+import functools
 import cv2
 import os
 
@@ -32,19 +33,40 @@ def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=1.0, threshold=0):
     return sharpened
 
                      
-def style_transfer(content_img,style_image, style_weight = 1, content_weight = 1, style_blur=False):
-    content_img = unsharp_mask(content_img,amount=1)
-    content_img = tf.image.resize(tf.convert_to_tensor(content_img,tf.float32)[tf.newaxis,...] / 255.,(512,512),preserve_aspect_ratio=True)
-    style_img = tf.convert_to_tensor(style_image,tf.float32)[tf.newaxis,...] / 255.
+def style_transfer(content_img, style_image, style_weight=1, content_weight=1, style_blur=False):
+    # Resize and preprocess the content image
+    content_img = unsharp_mask(content_img, amount=1)
+    content_img = tf.image.resize(
+        tf.convert_to_tensor(content_img, dtype=tf.float32)[tf.newaxis, ...] / 255.0,
+        (512, 512),
+        preserve_aspect_ratio=True
+    )
+    
+    # Resize and preprocess the style image
+    style_image = Image.fromarray(style_image).resize((256, 256))
+    style_img = tf.convert_to_tensor(np.array(style_image), dtype=tf.float32)[tf.newaxis, ...] / 255.0
+    
     if style_blur:
-        style_img=  tf.nn.avg_pool(style_img, [3,3], [1,1], "VALID")
-    style_img = tf.image.adjust_contrast(style_img, style_weight)        
-    content_img = tf.image.adjust_contrast(content_img,content_weight)     
-    content_img = tf.image.adjust_saturation(content_img, 2)        
-    content_img = tf.image.adjust_contrast(content_img,1.5)        
+        style_img = tf.nn.avg_pool(style_img, ksize=[3, 3], strides=[1, 1], padding="VALID")
+    
+    # Apply style weight to the style image
+    style_img = tf.image.adjust_contrast(style_img, style_weight)
+    
+    # Apply content weight and other adjustments to the content image
+    content_img = tf.image.adjust_contrast(content_img, content_weight)
+    content_img = tf.image.adjust_saturation(content_img, 2)
+    content_img = tf.image.adjust_contrast(content_img, 1.5)
+    
+    # Stylize the content image using the style image
     stylized_img = model(content_img, style_img)[0]
     
-    return Image.fromarray(np.uint8(stylized_img[0]*255))
+    # Convert the stylized image tensor to a NumPy array
+    stylized_img = tf.squeeze(stylized_img).numpy()
+    
+    # Convert the NumPy array to an image
+    stylized_img = np.clip(stylized_img * 255.0, 0, 255).astype(np.uint8)
+    
+    return Image.fromarray(stylized_img)
 
 
 
@@ -54,32 +76,30 @@ description = "Gradio Demo for Artistic Neural Style Transfer. To use it, simply
 article = "</br><p style='text-align: center'><a href='https://github.com/Mr-Hexi' target='_blank'>GitHub</a></p> "
 
 
-content_input = gr.inputs.Image(label="Upload an image to which you want the style to be applied.",)
-style_input = gr.inputs.Image( label="Upload Style Image ",shape= (256,256), )
-style_slider = gr.inputs.Slider(0,2,label="Adjust Style Density" ,default=1,)
-content_slider = gr.inputs.Slider(1,5,label="Content Sharpness" ,default=1,)
-# style_checkbox = gr.Checkbox(value=False,label="Tune Style(experimental)")
-
-
-examples  = [
-                ["Content/content_2.jpg","Styles/style_15.jpg",1.20,1.70,""], 
-                ["Content/content_4.jpg","Styles/Scream Edvard Munch.jpg",0.91,2.54,"style_checkbox"]
-            
-            ]
-interface = gr.Interface(fn=style_transfer,
-                         inputs=[content_input,
-                                style_input,
-                                style_slider ,
-                                content_slider,
-                                # style_checkbox
-                                ],
-                         outputs=gr.outputs.Image(),
-                         title=title,
-                         description=description,
-                         article=article,
-                         examples=examples,
-                         enable_queue=True
-                         )
-    
-    
-interface.launch()
+# Define inputs
+content_input = gr.Image(label="Upload an image to which you want the style to be applied.")
+style_input = gr.Image(label="Upload Style Image")  # Removed the shape parameter
+style_slider = gr.Slider(0, 2, label="Adjust Style Density", value=1)
+content_slider = gr.Slider(1, 5, label="Content Sharpness", value=1)
+style_checkbox = gr.Checkbox(value=False, label="Tune Style (experimental)")
+
+# Define examples
+examples = [
+    ["Content/content_2.jpg", "Styles/style_15.jpg", 1.20, 1.70, ""],
+    ["Content/content_4.jpg", "Styles/style_10.jpg", 0.91, 2.54, "style_checkbox"]
+]
+
+# Define the interface
+interface = gr.Interface(
+    fn=style_transfer,
+    inputs=[content_input, style_input, style_slider, content_slider, style_checkbox],
+    outputs=gr.Image(),
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    allow_flagging="never",
+)
+# Launch the interface
+interface.launch(debug=True)
+