fix

app.py

@@ -1,11 +1,11 @@
 import os, torch
-from model.BiSeNet.build_bisenet import BiSeNet
 import gradio as gr
 from utils.imageHandling import hfImageToTensor, preprocessing
+from model.modelLoading import loadModel
 
 # %% prediction on an image
 
-def predict(inputImage: torch.Tensor, model: BiSeNet) -> torch.Tensor:
+def predict(inputImage: torch.Tensor, model) -> torch.Tensor:
     """
     Predict the segmentation mask for the input image using the provided model.
     
@@ -22,45 +22,6 @@ def predict(inputImage: torch.Tensor, model: BiSeNet) -> torch.Tensor:
     return output[0].argmax(dim=0, keepdim=True).cpu()  
 
 
-
-# %% load model
-
-def loadModel(model:str = 'bisenet', device: str = 'cpu')->BiSeNet:
-    """
-    Load the specified model and move it to the given device.
-    
-    Args:
-        model (str): model to be loaded.
-        device (str): Device to load the model onto ('cpu' or 'cuda').
-        
-    Returns:
-        model (BiSeNet): The loaded BiSeNet model.
-    """
-    match model.lower() if isinstance(model, str) else model:
-        case 'bisenet': model = loadBiSeNet(device)
-        case _: raise NotImplementedError(f"Model {model} is not implemented. Please choose 'bisenet' .")
-    
-    return model
-
-
-# BiSeNet model loading function
-def loadBiSeNet(device: str = 'cpu') -> BiSeNet:
-    """
-    Load the BiSeNet model and move it to the specified device.
-    
-    Args:
-        device (str): Device to load the model onto ('cpu' or 'cuda').
-        
-    Returns:
-        model (BiSeNet): The loaded BiSeNet model.
-    """
-    model = BiSeNet(n_classes=19, context_path='resnet18').to(device)
-    model.load_state_dict(torch.load('./weights/BiSeNet/weightADV.pth', map_location=device))
-    model.eval()
-    
-    return model
-
-
 # %% Gradio interface
 def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -71,6 +32,7 @@ def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
 with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
     gr.Markdown("## 🧠 Image Segmentation with BiSeNet and BiSeNetV2")
     gr.Markdown("Upload an image and choose your preferred model for segmentation.")
+    gr.Markdown('A small user interface created to run semantic segmentation on images using city scapes like predictions and real time segmentation networks.')
 
     with gr.Row():
         with gr.Column():
@@ -89,5 +51,7 @@ with gr.Blocks(title="🔀 BiSeNet | BiSeNetV2 Predictor") as demo:
         inputs=[image_input, model_selector],
         outputs=[result_display]
     )
+    
+    gr.Markdown("Made by group 21 semantic segmentation project. ")
 
 demo.launch()

model/modelLoading.py

@@ -0,0 +1,41 @@
+import torch
+
+from model.BiSeNet.build_bisenet import BiSeNet
+
+
+# %% load model
+
+def loadModel(model:str = 'bisenet', device: str = 'cpu')->BiSeNet:
+    """
+    Load the specified model and move it to the given device.
+    
+    Args:
+        model (str): model to be loaded.
+        device (str): Device to load the model onto ('cpu' or 'cuda').
+        
+    Returns:
+        model (BiSeNet): The loaded BiSeNet model.
+    """
+    match model.lower() if isinstance(model, str) else model:
+        case 'bisenet': model = loadBiSeNet(device)
+        case _: raise NotImplementedError(f"Model {model} is not implemented. Please choose 'bisenet' .")
+    
+    return model
+
+
+# BiSeNet model loading function
+def loadBiSeNet(device: str = 'cpu') -> BiSeNet:
+    """
+    Load the BiSeNet model and move it to the specified device.
+    
+    Args:
+        device (str): Device to load the model onto ('cpu' or 'cuda').
+        
+    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.eval()
+    
+    return model