fixed weight loading

model/modelLoading.py

@@ -5,13 +5,13 @@ from model.BiSeNetV2.model import BiSeNetV2
 
 
 # BiSeNet model loading function
-def loadBiSeNet(device: str = 'cpu', weights:str='weightADV.pth') -> BiSeNet:
+def loadBiSeNet(device: str = 'cpu', weights:str='weight_Base.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
+        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.
@@ -23,18 +23,19 @@ def loadBiSeNet(device: str = 'cpu', weights:str='weightADV.pth') -> BiSeNet:
     return model
 
 
-def loadBiSeNetV2(device: str = 'cpu') -> BiSeNetV2:
+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').
+        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 (BiSeNetV2): The loaded BiSeNetV2 model.
     """
     model = BiSeNetV2(n_classes=19).to(device)
-    model.load_state_dict(torch.load('./weights/BiSeNetV2/weightADV.pth', map_location=device)['model_state_dict'])
+    model.load_state_dict(torch.load(f'./weights/BiSeNetV2/{weights}', map_location=device)['model_state_dict'])
     model.eval()
     
     return model