import gradio as gr
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

# Carregar o modelo RMBG-2.0
model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
model.to('cuda' if torch.cuda.is_available() else 'cpu')
model.eval()

# Função para remover o fundo da imagem
def remove_background(image):
    # Transformações necessárias
    transform = transforms.Compose([
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    # Aplicar transformações
    input_image = transform(image).unsqueeze(0).to('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Realizar a predição
    with torch.no_grad():
        output = model(input_image)[-1].sigmoid().cpu()
    
    # Processar a máscara
    mask = transforms.ToPILImage()(output[0].squeeze())
    mask = mask.resize(image.size)
    image.putalpha(mask)
    
    return image

    # Configurar a interface do Gradio
app = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Remoção de Fundo com BRIA AI 2.0"
)

# Executar o aplicativo
if __name__ == "__main__":
    app.launch(share=True)