import gradio as gr
import numpy as np
import torch
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

MODEL_ID = "mipedro1/segformer-buildings-segmentation"
processor = SegformerImageProcessor.from_pretrained(MODEL_ID)
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_ID)
model.eval()

def predict(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    upsampled = torch.nn.functional.interpolate(
        outputs.logits,
        size=(image.height, image.width),
        mode="bilinear",
        align_corners=False,
    )
    pred = upsampled.argmax(dim=1).squeeze().numpy()
    mask = (pred * 255).astype(np.uint8)
    return Image.fromarray(mask)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Ortofoto"),
    outputs=gr.Image(label="Edificios detectados"),
    title="Detección de edificios en ortofotos",
)

demo.launch()