alkzar90 commited on
Commit
3237429
1 Parent(s): aad7f4e

Change preprocessing function

Browse files
Files changed (2) hide show
  1. app.py +10 -5
  2. requirements.txt +1 -0
app.py CHANGED
@@ -3,15 +3,18 @@ import random
3
  import numpy as np
4
  import torch
5
  from torch import nn
6
- from transformers import (SegformerFeatureExtractor,
7
- SegformerForSemanticSegmentation)
8
 
9
 
10
  MODEL_PATH="./best_model_test/"
11
 
12
  device = torch.device("cpu")
13
 
14
- preprocessor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
 
 
15
  model = SegformerForSemanticSegmentation.from_pretrained(MODEL_PATH)
16
  model.eval()
17
 
@@ -43,8 +46,9 @@ def visualize_instance_seg_mask(mask):
43
  def query_image(img):
44
  """Función para generar predicciones a la escala origina"""
45
  inputs = preprocessor(images=img, return_tensors="pt")
 
46
  with torch.no_grad():
47
- preds = model(**inputs)["logits"]
48
  preds_upscale = upscale_logits(preds, preds.shape[2])
49
  predict_label = torch.argmax(preds_upscale, dim=1).to(device)
50
  result = predict_label[0,:,:].detach().cpu().numpy()
@@ -55,7 +59,8 @@ demo = gr.Interface(
55
  query_image,
56
  inputs=[gr.Image(type="pil")],
57
  outputs="image",
58
- title="SegFormer Model for rock glacier image segmentation"
 
59
  )
60
 
61
  demo.launch()
 
3
  import numpy as np
4
  import torch
5
  from torch import nn
6
+ from torchvision import transforms
7
+ from transformers import SegformerForSemanticSegmentation)
8
 
9
 
10
  MODEL_PATH="./best_model_test/"
11
 
12
  device = torch.device("cpu")
13
 
14
+ preprocessor = transforms.Compose([
15
+ transforms.resize(128),
16
+ transforms.ToTensor()
17
+ ])
18
  model = SegformerForSemanticSegmentation.from_pretrained(MODEL_PATH)
19
  model.eval()
20
 
 
46
  def query_image(img):
47
  """Función para generar predicciones a la escala origina"""
48
  inputs = preprocessor(images=img, return_tensors="pt")
49
+ inputs = preprocessor(img).unsqueeze(0)
50
  with torch.no_grad():
51
+ preds = model(inputs)["logits"]
52
  preds_upscale = upscale_logits(preds, preds.shape[2])
53
  predict_label = torch.argmax(preds_upscale, dim=1).to(device)
54
  result = predict_label[0,:,:].detach().cpu().numpy()
 
59
  query_image,
60
  inputs=[gr.Image(type="pil")],
61
  outputs="image",
62
+ title="Skyguard: segmentador de glaciares de roca 🛰️ +️ 🛡️ ️",
63
+ 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)"
64
  )
65
 
66
  demo.launch()
requirements.txt CHANGED
@@ -1,3 +1,4 @@
1
  torch
 
2
  transformers
3
  numpy
 
1
  torch
2
+ torchvision
3
  transformers
4
  numpy