Change preprocessing function

app.py

@@ -3,15 +3,18 @@ import random
 import numpy as np
 import torch
 from torch import nn
-from transformers import (SegformerFeatureExtractor,
-                          SegformerForSemanticSegmentation)
+from torchvision import transforms
+from transformers import SegformerForSemanticSegmentation)
 
 
 MODEL_PATH="./best_model_test/"
 
 device = torch.device("cpu")
 
-preprocessor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+preprocessor = transforms.Compose([
+                        transforms.resize(128),
+                        transforms.ToTensor()
+                 ])
 model = SegformerForSemanticSegmentation.from_pretrained(MODEL_PATH)
 model.eval()
 
@@ -43,8 +46,9 @@ def visualize_instance_seg_mask(mask):
 def query_image(img):
     """Función para generar predicciones a la escala origina"""
     inputs = preprocessor(images=img, return_tensors="pt")
+    inputs = preprocessor(img).unsqueeze(0)
     with torch.no_grad():
-        preds = model(**inputs)["logits"]
+        preds = model(inputs)["logits"]
     preds_upscale = upscale_logits(preds, preds.shape[2])
     predict_label = torch.argmax(preds_upscale, dim=1).to(device)
     result = predict_label[0,:,:].detach().cpu().numpy()
@@ -55,7 +59,8 @@ demo = gr.Interface(
     query_image,
     inputs=[gr.Image(type="pil")],
     outputs="image",
-    title="SegFormer Model for rock glacier image segmentation"
+    title="Skyguard: segmentador de glaciares de roca  🛰️ +️ 🛡️ ️",
+    description="Modelo de segmentación de imágenes para detectar glaciares de roca.<br> Se entrenó un modelo [nvidia/SegFormer](https://huggingface.co/nvidia/mit-b0) con _fine-tuning_ en el [rock-glacier-dataset](https://huggingface.co/datasets/alkzar90/rock-glacier-dataset)"
 )
 
 demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
 torch
+torchvision
 transformers
 numpy