simpler model loading

app.py

@@ -13,10 +13,10 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 MODELS = {
-    "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')
+    "BISENET-BASE": loadBiSeNet('bisenet', device, 'weight_Base.pth'),
+    "BISENET-BEST": loadBiSeNet('bisenet', device, 'weight_Best.pth'),
+    "BISENETV2-BASE": loadBiSeNetV2('bisenetv2', device, 'weight_Base.pth'),
+    "BISENETV2-BEST": loadBiSeNetV2('bisenetv2', device, 'weight_Best.pth')
 }
 
 image_list = loadPreloadedImages(gta_image_dir, city_image_dir, turin_image_dir)

model/modelLoading.py

@@ -4,38 +4,37 @@ from model.BiSeNet.build_bisenet import BiSeNet
 from model.BiSeNetV2.model import BiSeNetV2
 
 
-# BiSeNet model loading function
-def loadBiSeNet(device: str = 'cpu', weights:str='weight_Base.pth') -> BiSeNet:
+# Model loading function
+def loadModel(model:str = 'bisenet', device: str = 'cpu', weights:str='weight_Base.pth') -> BiSeNet | BiSeNetV2:
     """
-    Load the BiSeNet model and move it to the specified device.
+    Load the BiSeNet or BiSeNetV2 model and move it to the specified device.
+    This function supports loading different versions of the model based on the provided `model` argument.
+    The model weights are loaded from the specified `weights` file.
+    The model is set to evaluation mode after loading.
     
     Args:
+        model (str): The type of model to load. Options are 'bisenet', 'bisenet_base', 'bisenet_best', 'bisenetv2', 'bisenetv2_base', 'bisenetv2_best'.
+                    Default is 'bisenet'.
         device (str): Device to load the model onto ('cpu' or 'cuda'). Default is 'cpu'.
         weights (str): weights file to be loaded. Default is 'weight_Base.pth'.
         
     Returns:
-        model (BiSeNet): The loaded BiSeNet model.
+        model (BiSeNet | BiSeNetV2): The loaded BiSeNet or BiSeNetV2 model.
     """
-    model = BiSeNet(num_classes=19, context_path='resnet18').to(device)
-    model.load_state_dict(torch.load(f'./weights/BiSeNet/{weights}', map_location=device)['model_state_dict'])
-    model.eval()
     
-    return model
+    match model.lower() if isinstance(model, str) else model:
+        case 'bisenet' | 'bisenet_base' | 'bisenet_best': 
+            model = BiSeNet(num_classes=19, context_path='resnet18').to(device)
+            modelStateDict = torch.load(f'./weights/BiSeNet/{weights}', map_location=device)
+            model.load_state_dict(modelStateDict['model_state_dict'] if 'model_state_dict' in modelStateDict else modelStateDict)
 
+        case 'bisenetv2' | 'bisenetv2_base' | 'bisenetv2_best':
+            model = BiSeNetV2(n_classes=19).to(device)
+            modelStateDict = torch.load(f'./weights/BiSeNetV2/{weights}', map_location=device)
+            model.load_state_dict(modelStateDict['model_state_dict'] if 'model_state_dict' in modelStateDict else modelStateDict)
 
-def loadBiSeNetV2(device: str = 'cpu', weights:str='weight_Base.pth') -> BiSeNetV2:
-    """
-    Load the BiSeNetV2 model and move it to the specified device.
-    
-    Args:
-        device (str): Device to load the model onto ('cpu' or 'cuda'). Default is 'cpu'.
-        weights (str): weights file to be loaded. Default is 'weight_Base.pth'.
+        case _: raise NotImplementedError(f"Model {model} is not implemented. Please choose 'bisenet' or 'bisenetv2'.")
         
-    Returns:
-        model (BiSeNetV2): The loaded BiSeNetV2 model.
-    """
-    model = BiSeNetV2(n_classes=19).to(device)
-    model.load_state_dict(torch.load(f'./weights/BiSeNetV2/{weights}', map_location=device)['model_state_dict'])
     model.eval()
     
     return model