added legend

app.py

@@ -1,7 +1,7 @@
 import torch
 import gradio as gr
 
-from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing, loadPreloadedImages
+from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing, loadPreloadedImages, legendHandling
 from model.modelLoading import loadBiSeNet, loadBiSeNetV2
 
 
@@ -76,6 +76,15 @@ with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
         with gr.Column():
             result_display = gr.Image(label="Model prediction", visible=True)
             error_text = gr.Markdown("", visible=False)
+            with gr.Row():
+                gr.Markdown("The legend of the classes is the following (format **ID**: **name** - Color: **color**)")
+                string = ""
+                for i, (id, name, color) in enumerate(legendHandling()):
+                    if i and not i % 3:
+                        gr.Markdown(string)
+                        string = ""
+                    
+                    string += f"**{id}**: {name} - Color: {color}"
 
     with gr.Row():
         gr.Markdown("## Preloaded images to be used for testing the model")

utils/imageHandling.py

@@ -70,6 +70,24 @@ def print_mask(mask:torch.Tensor, numClasses:int=19)->None:
         new_mask[mask == i] = torch.tensor(colors[i][:3], dtype=torch.uint8)
     return new_mask.permute(2,0,1)  
 
+
+def legendHandling()->list[int, str, str]:
+    """
+    Returns a sorted list of tuples containing class IDs, names, and colors for semantic segmentation.
+    
+    Each tuple contains:
+        - Class ID (int)
+        - Class name (str)
+        - Color description (str)
+    The list is sorted by class ID.
+    """
+    return sorted([[0, "road", "dark purple"], [1, "sidewalk", "light purple / pink"], [2, "building", "dark gray"], [3, "wall", "blue + grey"], 
+            [4, "fence", "beige"], [5, "pole", "grey"], [6, "traffic light", "orange"], [7, "traffic sign", "yellow"], [8, "vegetation", "dark green"], 
+            [9, "terrain", "light green"], [10, "sky", "blue"], [11, "person", "dark red"], [12, "rider", "light red"], [13, "car", "blue"], 
+            [14, "truck", "dark blue"], [15, "bus", "dark blue"], [16, "train", "blue + green"], [17, "motorcycle", "light blue"], [18, "bicycle", "velvet"]
+    ], key=lambda x: x[0])
+
+
 # %% postprocessing
 def postprocessing(pred: torch.Tensor) -> torch.Tensor:
     """