Update app.py

app.py

@@ -6,10 +6,7 @@ from PIL import Image
 from rembg import remove
 
 # Load the neural style transfer model from TensorFlow Hub
-style_transfer_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
-
-# Load a super-resolution model from TensorFlow Hub
-super_res_model = hub.load('https://tfhub.dev/captain-pool/esrgan-tf2/1')
+model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
 
 # Function to convert tensor to image
 def tensor_to_image(tensor):
@@ -22,11 +19,12 @@ def tensor_to_image(tensor):
 
 # 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'))
+    input_rgb = np.array(image.convert('RGB'))  # Ensure RGB format
     output_rgba = np.array(output_image)
 
     alpha = output_rgba[:, :, 3]
@@ -43,54 +41,31 @@ 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 = style_transfer_model(tf.constant(content_image), tf.constant(style_image))
+    outputs = model(tf.constant(content_image), tf.constant(style_image))
     stylized_image = outputs[0]
 
     return tensor_to_image(stylized_image)
 
-# Super-resolution function
-def apply_super_resolution(image):
-    image = np.array(image).astype(np.float32) / 255.0
-    sr_image = super_res_model(tf.constant(image[np.newaxis, ...]))
-    sr_image = sr_image[0]
-    return tensor_to_image(sr_image)
-
 # Function to process image
 def process_image(content_image, style_image):
-    # Ensure style_image is a PIL Image
-    if isinstance(style_image, np.ndarray):
-        style_image = Image.fromarray(style_image)
-
     foreground, background = separate_foreground_background(content_image)
 
     # Convert to RGB format by removing the alpha channel
     foreground_rgb = np.array(foreground.convert('RGB'))
     background_rgb = np.array(background)
 
-    styled_foreground = apply_style_transfer(foreground_rgb, np.array(style_image.convert('RGB')), intensity=1.0)
-    styled_background = apply_style_transfer(background_rgb, np.array(style_image.convert('RGB')), intensity=0.3)
-
-    # Resize the styled images to match each other using a fixed target size
-    target_size = (512, 512)
-    styled_foreground = styled_foreground.resize(target_size, Image.LANCZOS)
-    styled_background = styled_background.resize(target_size, Image.LANCZOS)
-
-    # Apply super-resolution to enhance quality
-    enhanced_foreground = apply_super_resolution(styled_foreground)
-    enhanced_background = apply_super_resolution(styled_background)
-
-    enhanced_foreground_np = np.array(enhanced_foreground)
-    enhanced_background_np = np.array(enhanced_background)
+    styled_foreground = apply_style_transfer(foreground_rgb, style_image, intensity=1.0)
+    styled_background = apply_style_transfer(background_rgb, style_image, intensity=0.3)
 
-    # Resize the alpha channel to match the enhanced image size
-    alpha = np.array(foreground.resize(target_size))[:, :, 3] / 255.0
-    alpha = Image.fromarray((alpha * 255).astype(np.uint8)).resize(enhanced_foreground_np.shape[1::-1], Image.LANCZOS)
-    alpha = np.array(alpha) / 255.0
+    styled_foreground_np = np.array(styled_foreground)
+    styled_background_np = np.array(styled_background)
 
-    # Combine the images
-    combined_image_np = (enhanced_foreground_np * alpha[..., np.newaxis] +
-                         enhanced_background_np * (1 - alpha[..., np.newaxis]))
+    # Extract the alpha channel from the foreground
+    alpha = np.array(foreground)[:, :, 3] / 255.0
+    combined_image_np = (styled_foreground_np * alpha[..., np.newaxis] +
+                         styled_background_np * (1 - alpha[..., np.newaxis]))
 
     combined_image = Image.fromarray(np.clip(combined_image_np, 0, 255).astype(np.uint8))
 
@@ -104,5 +79,6 @@ gr.Interface(
     fn=process_image,
     inputs=[image1, image2],
     outputs=stylizedimg,
-    title='Stylized Foreground and Background Combination with Super-Resolution',
+    title='Stylized Foreground and Background Combination',
+    
 ).launch()