lambdas

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from PIL import Image
 
 from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing, loadPreloadedImages
-from model.modelLoading import loadModel
+from model.modelLoading import loadBiSeNet, loadBiSeNetV2
 
 
 ## %% CONSTANTS
@@ -12,7 +12,16 @@ city_image_dir = "./preloadedImages/cityScapes"
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
+MODELS = {
+    "BiSeNet": loadBiSeNet(device),
+    "BiSeNetV2": loadBiSeNetV2(device)
+}
+
+
 def load_example(index:int=0, useGta:bool=True)-> gr.Image:
+    print(type(index))
+    print(index)
+    print(useGta)
     example_img = loadPreloadedImages(gta_image_dir if useGta else city_image_dir)
     return gr.update(value=example_img[index if 0 <= index < len(example_img) else 0], visible=True)
 
@@ -44,10 +53,12 @@ def run_prediction(image: gr.Image, selected_model: str)-> tuple[torch.Tensor]:
     if selected_model is None:
         return (gr.update(value=None, visible=False),  gr.update(value=f"❌ No model selected for prediction.", visible=True))
     
+    if not isinstance(selected_model, str) or selected_model.strip().upper() not in MODELS:
+        return (gr.update(value=None, visible=False),  gr.update(value=f"❌ Invalid model selected.", visible=True))
+
     try:
-        model = loadModel(selected_model, device) 
         image = hfImageToTensor(image, width=1024, height=512)
-        prediction = predict(image, model)
+        prediction = predict(image, MODELS[selected_model.strip().upper()])
         prediction = postprocessing(prediction)
     except Exception as e:
         return (gr.update(value=None, visible=False),  gr.update(value=f"❌ {str(e)}.", visible=True))
@@ -88,8 +99,8 @@ with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
             rows=1, height=256, allow_preview=False
     )
     
-    gta_gallery.select(fn=load_example, inputs=[gta_gallery, True], outputs=[image_input])
-    city_gallery.select(fn=load_example, inputs=[city_gallery, False], outputs=[image_input])
+    gta_gallery.select(fn=lambda x: load_example(x, True), inputs=[gta_gallery], outputs=[image_input])
+    city_gallery.select(fn=lambda x: load_example(x, False), inputs=[city_gallery], outputs=[image_input])
 
     submit_btn.click(
         fn=run_prediction,

model/modelLoading.py

@@ -3,25 +3,6 @@ import torch
 from model.BiSeNet.build_bisenet import BiSeNet
 from model.BiSeNetV2.model import BiSeNetV2
 
-# general loading function
-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 'bisenetv2': model = loadBiSeNetV2(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: