add mask to rgb label2color image

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import random
+import numpy as np
 import torch
 from torch import nn
 from transformers import (SegformerFeatureExtractor,
@@ -23,23 +25,35 @@ def upscale_logits(logit_outputs, size):
        align_corners=False
    )
 
+
+def visualize_instance_seg_mask(mask):
+    """Agrega colores RGB a cada una de las clases en la mask"""
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    labels = np.unique(mask)
+    label2color = {label: (random.randint(0, 1), 
+                           random.randint(0, 255), 	
+                           random.randint(0, 255)) for label in labels}
+    for i in range(image.shape[0]):
+        for j in range(image.shape[1]):
+            image[i, j, :] = label2color[mask[i, j]]
+        image = image / 255
+    return image
+
+
 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(**inputs)["logits"]
-    preds_upscale = upscale_logits(preds, img.shape[2])
+    preds_upscale = upscale_logits(preds, preds.shape[2])
     predict_label = torch.argmax(preds_upscale, dim=1).to(device)
-    return predict_label[0,:,:].detach().cpu().numpy()
-
+    result = predict_label[0,:,:].detach().cpu().numpy()
+    return visualize_instance_seg_mask(result)
 
-def visualize_instance_seg_mask(mask):
-    return mask
 
 demo = gr.Interface(
     query_image,
-    inputs=[gr.Image()],
+    inputs=[gr.Image(type="pil")],
     outputs="image",
     title="SegFormer Model for rock glacier image segmentation"
 )

requirements.txt

@@ -1,2 +1,3 @@
 torch
 transformers
+numpy