added new features to the Gradio interface and fixed some bugs in the prediction function.

app.py

@@ -39,7 +39,7 @@ def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
     return (gr.update(value=prediction, visible=True), gr.update(value="", visible=False))
 
 # Gradio UI
-with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
+with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
     gr.Markdown("## Semantic Segmentation with Real-Time Networks")
     gr.Markdown('A small user interface created to run semantic segmentation on images using city scapes like predictions and real time segmentation networks.')
     gr.Markdown("Upload an image and choose your preferred model for segmentation.")
@@ -54,16 +54,15 @@ with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
             image_input = gr.Image(type="pil", label="Upload image")
             submit_btn = gr.Button("Run prediction")
         with gr.Column():
-            result_display = gr.Image(label="Model prediction")
-            
-    error_text = gr.Textbox(label="Message", interactive=False, visible=False)
-    
+            result_display = gr.Image(label="Model prediction", visible=False)
+            error_text = gr.Markdown("", visible=False)
+
     submit_btn.click(
         fn=run_prediction,
         inputs=[image_input, model_selector],
         outputs=[result_display, error_text],
     )
-    
+
     gr.Markdown("Made by group 21 semantic segmentation project. ")
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'

utils/imageHandling.py

@@ -31,6 +31,43 @@ def preprocessing(image_tensor: torch.Tensor) -> torch.Tensor:
     return torchvision.transforms.functional.normalize(
         image_tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
     ).unsqueeze(0)
+    
+# %% print mask on a sem seg style
+def print_mask(mask:torch.Tensor, numClasses:int=19)->None:
+    """
+        Visualizes the segmentation mask by mapping each class to a specific color.
+    
+    Args:
+        mask (torch.Tensor): The segmentation mask to visualize.
+        numClasses (int, optional): Number of classes in the segmentation mask. Defaults to 19.
+    """
+    colors = [
+        (128, 64, 128),  # 0: road
+        (244, 35, 232),  # 1: sidewalk
+        (70, 70, 70),    # 2: building
+        (102, 102, 156), # 3: wall
+        (190, 153, 153), # 4: fence
+        (153, 153, 153), # 5: pole
+        (250, 170, 30),  # 6: traffic light
+        (220, 220, 0),   # 7: traffic sign
+        (107, 142, 35),  # 8: vegetation
+        (152, 251, 152), # 9: terrain
+        (70, 130, 180),  # 10: sky
+        (220, 20, 60),   # 11: person
+        (255, 0, 0),     # 12: rider
+        (0, 0, 142),     # 13: car
+        (0, 0, 70),      # 14: truck
+        (0, 60, 100),    # 15: bus
+        (0, 80, 100),    # 16: train
+        (0, 0, 230),     # 17: motorcycle
+        (119, 11, 32)    # 18: bicycle
+    ]
+
+    new_mask = torch.zeros((mask.shape[0], mask.shape[1], 3),dtype=torch.uint8)
+    new_mask[mask == 255] = (0,0,0)
+    for i in range (numClasses):
+        new_mask[mask == i] = colors[i][:3]
+    return new_mask.permute(2,0,1)  
 
 # %% postprocessing
 def postprocessing(pred: torch.Tensor) -> torch.Tensor:
@@ -43,4 +80,4 @@ def postprocessing(pred: torch.Tensor) -> torch.Tensor:
     Returns:
         torch.Tensor: Processed tensor of shape (3, H, W) for visualization.
     """ 
-    return torchvision.transforms.functional.to_pil_image(pred.squeeze(0).cpu().to(torch.uint8))
+    return torchvision.transforms.functional.to_pil_image(print_mask(pred.squeeze(0).cpu().to(torch.uint8)))

utils2.py

@@ -1,37 +0,0 @@
-import torch
-
-def print_mask(mask:torch.Tensor, numClasses:int=19)->None:
-    """
-        Visualizes the segmentation mask by mapping each class to a specific color.
-    
-    Args:
-        mask (torch.Tensor): The segmentation mask to visualize.
-        numClasses (int, optional): Number of classes in the segmentation mask. Defaults to 19.
-    """
-    colors = [
-        (128, 64, 128),  # 0: road
-        (244, 35, 232),  # 1: sidewalk
-        (70, 70, 70),    # 2: building
-        (102, 102, 156), # 3: wall
-        (190, 153, 153), # 4: fence
-        (153, 153, 153), # 5: pole
-        (250, 170, 30),  # 6: traffic light
-        (220, 220, 0),   # 7: traffic sign
-        (107, 142, 35),  # 8: vegetation
-        (152, 251, 152), # 9: terrain
-        (70, 130, 180),  # 10: sky
-        (220, 20, 60),   # 11: person
-        (255, 0, 0),     # 12: rider
-        (0, 0, 142),     # 13: car
-        (0, 0, 70),      # 14: truck
-        (0, 60, 100),    # 15: bus
-        (0, 80, 100),    # 16: train
-        (0, 0, 230),     # 17: motorcycle
-        (119, 11, 32)    # 18: bicycle
-    ]
-
-    new_mask = torch.zeros((mask.shape[0], mask.shape[1], 3),dtype=torch.uint8)
-    new_mask[mask == 255] = (0,0,0)
-    for i in range (numClasses):
-        new_mask[mask == i] = colors[i][:3]
-    return new_mask.permute(2,0,1)