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