added more weights

app.py

@@ -13,8 +13,10 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 MODELS = {
-    "BISENET": loadBiSeNet(device),
-    "BISENETV2": loadBiSeNetV2(device)
+    "BISENET-BASE": loadBiSeNet(device, 'weight_Base.pth'),
+    "BISENET-BEST": loadBiSeNet(device, 'weight_Best.pth'),
+    "BISENETV2-BASE": loadBiSeNetV2(device, 'weight_Base.pth'),
+    "BISENETV2-BEST": loadBiSeNetV2(device, 'weight_Best.pth')
 }
 
 image_list = loadPreloadedImages(gta_image_dir, city_image_dir, turin_image_dir)
@@ -68,8 +70,8 @@ with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload image")
             model_selector = gr.Radio(
-                choices=["BiSeNet", "BiSeNetV2"],
-                value="BiSeNet",
+                choices=["BiSeNet-base", "BiSeNet-Best", "BiSeNetV2-base", "BiSeNetV2-Best"],
+                value="BiSeNet-base",
                 label="Select the real time segmentation model"
                 )
             submit_btn = gr.Button("Run prediction")
@@ -82,10 +84,14 @@ with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
                 for i in range(0, len(legend), 2):
                     with gr.Row():
                         with gr.Column(scale=1):
-                            gr.Markdown(f"**{legend[i][1]}** → {legend[i][2]}")
+                            color_box0 = f"""<span style='display:inline-block; width:15px; height:15px; 
+                                background-color:rgb({legend[i][1][0]},{legend[i][1][1]},{legend[i][1][2]}); margin-left:6px; border:1px solid #000;'></span>"""
+                            gr.Markdown(f"**{legend[i][1]}** → {legend[i][2]} {color_box0}")
                         with gr.Column(scale=1):
                             if i + 1 < len(legend):
-                                gr.Markdown(f"**{legend[i+1][1]}** → {legend[i+1][2]}")
+                                color_box1 = f"""<span style='display:inline-block; width:15px; height:15px; 
+                                    background-color:rgb({legend[i+1][1][0]},{legend[i+1][1][1]},{legend[i+1][1][2]}); margin-left:6px; border:1px solid #000;'></span>"""
+                                gr.Markdown(f"**{legend[i+1][1]}** → {legend[i+1][2]} {color_box1}")
                             else:
                                 gr.Markdown("")  # Keeps spacing consistent if list is odd
 

model/modelLoading.py

@@ -5,18 +5,19 @@ from model.BiSeNetV2.model import BiSeNetV2
 
 
 # BiSeNet model loading function
-def loadBiSeNet(device: str = 'cpu') -> BiSeNet:
+def loadBiSeNet(device: str = 'cpu', weights:str='weightADV.pth') -> BiSeNet:
     """
     Load the BiSeNet model and move it to the specified device.
     
     Args:
         device (str): Device to load the model onto ('cpu' or 'cuda').
+        weights (str): weights file to be loaded
         
     Returns:
         model (BiSeNet): The loaded BiSeNet model.
     """
     model = BiSeNet(num_classes=19, context_path='resnet18').to(device)
-    model.load_state_dict(torch.load('./weights/BiSeNet/weightADV.pth', map_location=device)['model_state_dict'])
+    model.load_state_dict(torch.load(f'./weights/BiSeNet/{weights}', map_location=device)['model_state_dict'])
     model.eval()
     
     return model

utils/imageHandling.py

@@ -81,10 +81,11 @@ def legendHandling()->list[int, str, 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"]
+    return sorted([[0, "road", "dark purple", (128, 64, 128)], [1, "sidewalk", "light purple / pink", (244, 35, 232)], [2, "building", "dark gray", (70, 70, 70)], 
+                    [3, "wall", "blue + grey", (102, 102, 156)], [4, "fence", "beige", (190, 153, 153)], [5, "pole", "grey", (153, 153, 153)], [6, "traffic light", "orange", (250, 170, 30)], 
+                    [7, "traffic sign", "yellow", (220, 220, 0)], [8, "vegetation", "dark green", (107, 142, 35)], [9, "terrain", "light green", (152, 251, 152)], [10, "sky", "blue", (70, 130, 180)], 
+                    [11, "person", "dark red", (220, 20, 60)], [12, "rider", "light red", (255, 0, 0)], [13, "car", "blue", (0, 0, 142)], [14, "truck", "dark blue", (0, 0, 70)],
+                    [15, "bus", "dark blue", (0, 60, 100)], [16, "train", "blue + green", (0, 80, 100)], [17, "motorcycle", "light blue", (0, 0, 230)], [18, "bicycle", "velvet", (119, 11, 32)]
     ], key=lambda x: x[0])
 
 

weights/BiSeNet/weight_Base.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e81645b4453ec62eda383d16e9cf279607181975dd648fdd354c78b050e7c8ca
+size 50456242

weights/BiSeNet/{weightADV.pth → weight_Best.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:880db4160f20c87aecc13845ad691b1963fbce3d713b1dda1964457b9e0d8f0a
+oid sha256:372b7ecbd6c9e0fc6909370a58e3445552f3379db884ca984cd94d3de0cac66c
 size 121015606

weights/BiSeNetV2/weight_Base.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ff01a531a19ec5d63837d48c4f8157d8b6fa0e04960035bdd26a40c39cb5b67
+size 21146508