import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
import gradio as gr

# Carrega o modelo de transferência de estilo pré-treinado
style_transfer_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')

def load_image(image):
    # Função para processar a imagem para o modelo
    image = image.astype(np.float32)[np.newaxis, ...] / 255.
    if image.shape[-1] == 4:
        image = image[..., :3]
    return image

def style_transfer(content_image, style_image):
    # Processa as imagens
    content_image = load_image(content_image)
    style_image = load_image(style_image)

    # Executa a transferência de estilo
    stylized_image = style_transfer_model(tf.constant(content_image), tf.constant(style_image))[0]

    # Converte a imagem resultante para o formato correto
    stylized_image = np.array(stylized_image * 255, np.uint8)

    # Remove a dimensão do batch
    stylized_image = np.squeeze(stylized_image)
    return stylized_image

iface = gr.Interface(
    fn=style_transfer,
    inputs=["image", "image"],
    outputs="image"
)

iface.launch()