import gradio as gr import tensorflow as tf import tensorflow_hub as hub import numpy as np from PIL import Image from rembg import remove # Load the neural style transfer model from TensorFlow Hub 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): tensor = tensor * 255 tensor = np.array(tensor, dtype=np.uint8) if np.ndim(tensor) > 3: assert tensor.shape[0] == 1 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 output_rgba = np.array(output_image) alpha = output_rgba[:, :, 3] alpha3 = np.dstack((alpha, alpha, alpha)) background_rgb = input_rgb.astype(np.float32) * (1 - alpha3.astype(np.float32) / 255) background_rgb = background_rgb.astype(np.uint8) foreground = Image.fromarray(output_rgba) 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)) stylized_image = outputs[0] return tensor_to_image(stylized_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): style_image = Image.fromarray(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 combined_image_np = (styled_foreground_np * alpha_resized[..., np.newaxis] + styled_background_np * (1 - alpha_resized[..., np.newaxis])) combined_image = Image.fromarray(np.clip(combined_image_np, 0, 255).astype(np.uint8)) return combined_image # Gradio interface setup image1 = gr.Image(label="Content Image") image2 = gr.Image(label="Style Image") stylizedimg = gr.Image(label="Result") gr.Interface( fn=process_image, inputs=[image1, image2], outputs=stylizedimg, title='Stylized Foreground and Background Combination', ).launch()