Mainapp.py

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()