Update app.py

app.py

@@ -1,15 +1,47 @@
+import os
+import cv2
 import gradio as gr
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
 import tensorflow as tf
 import tensorflow_hub as hub
-import numpy as np
-from PIL import Image, ImageEnhance, ImageFilter
 from rembg import remove
-import cv2
 
 # Load the neural style transfer model from TensorFlow Hub
-model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
+style_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
+
+# Set up ONNX runtime for cartoonizer model
+_sess_options = ort.SessionOptions()
+_sess_options.intra_op_num_threads = os.cpu_count()
+MODEL_SESS = ort.InferenceSession(
+    "cartoonizer.onnx", _sess_options, providers=["CPUExecutionProvider"]
+)
+
+def preprocess_image(image: Image) -> np.ndarray:
+    image = np.array(image)
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    h, w, c = np.shape(image)
+    if min(h, w) > 720:
+        if h > w:
+            h, w = int(720 * h / w), 720
+        else:
+            h, w = 720, int(720 * w / h)
+    image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
+    h, w = (h // 8) * 8, (w // 8) * 8
+    image = image[:h, :w, :]
+    image = image.astype(np.float32) / 127.5 - 1
+    return np.expand_dims(image, axis=0)
+
+def cartoonize(image: np.ndarray) -> Image:
+    image = preprocess_image(image)
+    results = MODEL_SESS.run(None, {"input_photo:0": image})
+    output = (np.squeeze(results[0]) + 1.0) * 127.5
+    output = np.clip(output, 0, 255).astype(np.uint8)
+    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(output)
 
-# Function to convert tensor to image
 def tensor_to_image(tensor):
     tensor = tensor * 255
     tensor = np.array(tensor, dtype=np.uint8)
@@ -18,19 +50,16 @@ def tensor_to_image(tensor):
         tensor = tensor[0]
     return Image.fromarray(tensor)
 
-# Function to resize image while maintaining aspect ratio
 def resize_image(image, target_size):
     image.thumbnail(target_size, Image.LANCZOS)
     return image
 
-# Function to separate foreground and background
 def separate_foreground_background(image):
-    # Ensure the image is a PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     
     output_image = remove(image)
-    input_rgb = np.array(image.convert('RGB'))  # Ensure RGB format
+    input_rgb = np.array(image.convert('RGB'))
     output_rgba = np.array(output_image)
 
     alpha = output_rgba[:, :, 3]
@@ -42,67 +71,17 @@ def separate_foreground_background(image):
     background = Image.fromarray(background_rgb)
     return foreground, background
 
-# Style transfer function
 def apply_style_transfer(content_image, style_image, intensity=1.0):
     content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
     style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
 
-    # Adjust the style intensity
     style_image = style_image * intensity
-    outputs = model(tf.constant(content_image), tf.constant(style_image))
+    outputs = style_model(tf.constant(content_image), tf.constant(style_image))
     stylized_image = outputs[0]
 
     return tensor_to_image(stylized_image)
 
-# Function to apply glossy effect
-def apply_glossy_effect(image):
-    # Convert the input image to a PIL Image
-    image = Image.fromarray(image)
-
-    # Enhance color saturation
-    color_enhancer = ImageEnhance.Color(image)
-    vibrant_image = color_enhancer.enhance(1)  # Increase saturation
-
-    # Enhance brightness
-    brightness_enhancer = ImageEnhance.Brightness(vibrant_image)
-    bright_image = brightness_enhancer.enhance(1)  # Adjust brightness
-
-    # Apply a soft blur to create a subtle glow
-    glow_image = bright_image.filter(ImageFilter.GaussianBlur(radius=2))
-
-    # Blend the original and glow images to create the glossy effect
-    glossy_image = Image.blend(bright_image, glow_image, alpha=0.9)
-
-    # Enhance contrast
-    contrast_enhancer = ImageEnhance.Contrast(glossy_image)
-    glossy_image = contrast_enhancer.enhance(1.5)  # Increase contrast
-
-    # Enhance sharpness
-    sharpness_enhancer = ImageEnhance.Sharpness(glossy_image)
-    glossy_image = sharpness_enhancer.enhance(2.8)  # Increase sharpness
-
-    # Apply edge enhancement
-    edges = glossy_image.filter(ImageFilter.FIND_EDGES)
-    glossy_image = Image.blend(glossy_image, edges, alpha=0.1)  # Subtle edge enhancement
-
-    # Apply vignette effect
-    np_img = np.array(glossy_image)
-    rows, cols, _ = np_img.shape
-    kernel_x = cv2.getGaussianKernel(cols, cols / 2)
-    kernel_y = cv2.getGaussianKernel(rows, rows / 2)
-    kernel = kernel_y @ kernel_x.T
-    mask = 255 * kernel / np.max(kernel)
-    vignette = np.zeros_like(np_img, dtype=np.uint8)
-    for i in range(3):  # Apply vignette on each channel
-        vignette[..., i] = np_img[..., i] * mask.astype(np.float32) / 255
-    glossy_image = Image.fromarray(np.clip(vignette, 0, 255).astype(np.uint8))
-
-    # Convert the glossy image back to a NumPy array
-    return np.array(glossy_image)
-
-# Function to process image
 def process_image(content_image, style_image):
-    # Ensure content_image and style_image are PIL Images
     if isinstance(content_image, np.ndarray):
         content_image = Image.fromarray(content_image)
     if isinstance(style_image, np.ndarray):
@@ -110,29 +89,23 @@ def process_image(content_image, style_image):
 
     foreground, background = separate_foreground_background(content_image)
 
-    # Define the target size based on the content image
     target_size = content_image.size
 
-    # Resize the foreground and background to match the target size
     foreground = resize_image(foreground, target_size)
     background = resize_image(background, target_size)
 
-    # Convert to RGB format by removing the alpha channel
     foreground_rgb = np.array(foreground.convert('RGB'))
     background_rgb = np.array(background)
 
-    # Apply style transfer
     styled_foreground = apply_style_transfer(foreground_rgb, np.array(style_image), intensity=1.0)
     styled_background = apply_style_transfer(background_rgb, np.array(style_image), intensity=0.3)
 
-    # Ensure both styled images have the same size
     styled_foreground = styled_foreground.resize(target_size, Image.LANCZOS)
     styled_background = styled_background.resize(target_size, Image.LANCZOS)
 
     styled_foreground_np = np.array(styled_foreground)
     styled_background_np = np.array(styled_background)
 
-    # Extract the alpha channel from the foreground
     alpha = np.array(foreground)[:, :, 3] / 255.0
     alpha_resized = np.array(foreground.resize(target_size))[:, :, 3] / 255.0
 
@@ -141,8 +114,8 @@ def process_image(content_image, style_image):
 
     combined_image = Image.fromarray(np.clip(combined_image_np, 0, 255).astype(np.uint8))
 
-    # Apply the glossy effect to the combined image
-    final_image = apply_glossy_effect(np.array(combined_image))
+    # Apply cartoonization to the final combined image
+    final_image = cartoonize(np.array(combined_image))
 
     return final_image
 
@@ -154,5 +127,5 @@ gr.Interface(
     fn=process_image,
     inputs=[image1, image2],
     outputs=stylizedimg,
-    
+    title='Stylized Foreground and Background Combination with Cartoon Effect',
 ).launch()