File size: 1,130 Bytes
bf601e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import gradio as gr
import torch
from transformers import (SegformerFeatureExtractor,
SegformerForSemanticSegmentation)
from PIL import Image
MODEL_PATH="./best_model_test/"
device = torch.device("cpu")
preprocessor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained(MODEL_PATH)
model.eval()
def query_image(img):
"""Función para generar predicciones a la escala origina"""
inputs = preprocessor(images=img, return_tensors="pt")
with torch.no_grad():
preds = model(inputs.unsqueeze(0).to(device))["logits"]
#preds = model(image.unsqueeze(0).to(device))["logits"]
preds_upscale = upscale_logits_modified(preds, image.shape[2])
predict_label = torch.argmax(preds_upscale, dim=1).to(device)
return predict_label[0,:,:].detach().cpu().numpy()
def visualize_instance_seg_mask(mask):
return mask
demo = gr.Interface(
query_image,
inputs=[gr.Image()],
outputs="image",
title="SegFormer Model for rock glacier image segmentation"
)
demo.launch()
|