Add multiple improvements

src/app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 from utils import load_specific_model, inference
-import markdown
 
-current_model = None  # Initialize the current model as None
+# current_model = None  # Initialize the current model as None
+MODEL_NAMES = ["EfficientNet-B3", "EfficientNet-B4", "vgg19", "resnet50", "dinov2_vits14"]
 
 # Define a set of example images
 example_images = [
@@ -13,23 +13,17 @@ example_images = [
     ("Beispielbild Wertstoff", "src/examples/Wertstoff.jpg")
 ]
 
-def load_model(model_name):
-    global current_model
-    if model_name is None:
+def predict(inp_image, inp_dropdown):
+    if inp_dropdown is None:
         raise gr.Error("No model selected!")
-    if current_model is not None:
-        current_model = None
-
-    current_model = load_specific_model(model_name)
-    current_model.eval()
-
-def predict(inp):
-    global current_model
-    print(f"\nPredicting with {current_model}\n")
-    if current_model is None:
-        raise gr.Error("No model loaded!")
-
-    confidences = inference(current_model, inp)
+    if inp_image is None:
+        raise gr.Error("No image uploaded!")
+    if inp_dropdown not in MODEL_NAMES:
+        raise gr.Error("Invalid model selected!")
+    
+    print(f"\nInput: {inp_dropdown}\n")
+    current_model = load_specific_model(inp_dropdown)
+    confidences = inference(current_model, inp_image)
     print(f"\nConfidences: {confidences}\n")
     return confidences
 
@@ -57,15 +51,15 @@ with gr.Blocks() as demo:
             )
 
             model_dropdown = gr.Dropdown(
-                ["EfficientNet-B3", "EfficientNet-B4", "vgg19", "resnet50", "dinov2_vits14"],
+                MODEL_NAMES,
+                value="EfficientNet-B3",
                 label="Model",
                 info="Select a model to use.",
                 scale=1,
             )
-            model_dropdown.change(load_model, model_dropdown, show_progress=True, queue=True)
             
-            predict_button = gr.Button(label="Predict", info="Click to make a prediction.", scale=1)
-            predict_button.click(fn=predict, inputs=user_image, outputs=output, queue=True)
+            predict_button = gr.Button(value="Predict", label="Predict", info="Click to make a prediction.", scale=1)
+            predict_button.click(fn=predict, inputs=[user_image, model_dropdown], outputs=output, queue=True)
 
     gr.Markdown("## Example Images")
     gr.Markdown("You can just drag and drop these images into the image uploader above!")

src/utils.py

@@ -3,6 +3,7 @@ from torchvision import models
 from efficientnet_pytorch import EfficientNet
 import torch
 from CustomModels import DinoVisionClassifier
+from functools import lru_cache
 
 classes = {0: 'Glas', 1: 'Organic', 2: 'Papier', 3: 'Restmüll', 4: 'Wertstoff'}
 
@@ -23,6 +24,7 @@ transform_dinov2 = transforms.Compose(
      ]
 )
 
+@lru_cache
 def load_specific_model(model_name):
     current_model = None
     if model_name == "EfficientNet-B3":