simpler model loading

app.py

@@ -2,7 +2,7 @@ import torch
 import gradio as gr
 
 from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing, loadPreloadedImages, legendHandling
-from model.modelLoading import loadBiSeNet, loadBiSeNetV2
+from model.modelLoading import loadModel
 
 
 ## %% CONSTANTS
@@ -13,10 +13,10 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 MODELS = {
-    "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')
+    "BISENET-BASE": loadModel('bisenet', device, 'weight_Base.pth'),
+    "BISENET-BEST": loadModel('bisenet', device, 'weight_Best.pth'),
+    "BISENETV2-BASE": loadModel('bisenetv2', device, 'weight_Base.pth'),
+    "BISENETV2-BEST": loadModel('bisenetv2', device, 'weight_Best.pth')
 }
 
 image_list = loadPreloadedImages(gta_image_dir, city_image_dir, turin_image_dir)