2025-07-31 21:34 🐛

app.py

@@ -22,6 +22,7 @@ alpha = 0.8
 EPS = 1e-8
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 loaded_model = None
+current_model_config = {"variant": None, "dataset": None, "metric": None}
 
 pretrained_models = [
     "ZIP-B @ ShanghaiTech A", "ZIP-B @ ShanghaiTech B", "ZIP-B @ UCF-QNRF", "ZIP-B @ NWPU-Crowd",
@@ -31,6 +32,39 @@ pretrained_models = [
     "ZIP-P @ ShanghaiTech A", "ZIP-P @ ShanghaiTech B", "ZIP-P @ UCF-QNRF"
 ]
 
+# -----------------------------
+# Model management functions
+# -----------------------------
+def update_model_if_needed(variant_dataset: str, metric: str):
+    """
+    Load a new model only if the configuration has changed.
+    """
+    global loaded_model, current_model_config
+    variant, dataset = variant_dataset.split(" @ ")
+
+    if dataset == "ShanghaiTech A":
+        dataset_name = "sha"
+    elif dataset == "ShanghaiTech B":
+        dataset_name = "shb"
+    elif dataset == "UCF-QNRF":
+        dataset_name = "qnrf"
+    elif dataset == "NWPU-Crowd":
+        dataset_name = "nwpu"
+
+    if (loaded_model is None or 
+        current_model_config["variant"] != variant or 
+        current_model_config["dataset"] != dataset_name or 
+        current_model_config["metric"] != metric):
+        
+        print(f"Loading new model: {variant} @ {dataset} with {metric} metric")
+        loaded_model = load_model(variant=variant, dataset=dataset_name, metric=metric)
+        current_model_config = {"variant": variant, "dataset": dataset_name, "metric": metric}
+        return f"Model loaded: {variant} @ {dataset} ({metric})"
+    else:
+        print(f"Using cached model: {variant} @ {dataset} with {metric} metric")
+        return f"Model already loaded: {variant} @ {dataset} ({metric})"
+
+
 # -----------------------------
 # Define the model architecture
 # -----------------------------
@@ -242,7 +276,7 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
     Given an input image, preprocess it, run the model to obtain a density map,
     compute the total crowd count, and prepare the density map for display.
     """
-    # global loaded_model
+    global loaded_model, current_model_config
     variant, dataset = variant_dataset.split(" @ ")
 
     if dataset == "ShanghaiTech A":
@@ -254,8 +288,16 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
     elif dataset == "NWPU-Crowd":
         dataset_name = "nwpu"
 
-    if loaded_model is None:
+    if (loaded_model is None or 
+        current_model_config["variant"] != variant or 
+        current_model_config["dataset"] != dataset_name or 
+        current_model_config["metric"] != metric):
+        
+        print(f"Loading new model: {variant} @ {dataset} with {metric} metric")
         loaded_model = load_model(variant=variant, dataset=dataset_name, metric=metric)
+        current_model_config = {"variant": variant, "dataset": dataset_name, "metric": metric}
+    else:
+        print(f"Using cached model: {variant} @ {dataset} with {metric} metric")
 
     if not hasattr(loaded_model, "input_size"):
         if dataset_name == "sha":
@@ -409,6 +451,12 @@ with gr.Blocks() as demo:
                 label="Select Best Metric"
             )
 
+            model_status = gr.Textbox(
+                label="Model Status",
+                value="No model loaded",
+                interactive=False
+            )
+
             input_img = gr.Image(label="Input Image", sources=["upload", "clipboard"], type="pil")
             submit_btn = gr.Button("Predict")
     
@@ -422,6 +470,21 @@ with gr.Blocks() as demo:
             output_sampling_zero_map = gr.Image(label="Sampling Zero Map", type="pil")
             output_complete_zero_map = gr.Image(label="Complete Zero Map", type="pil")
 
+    def on_model_change(variant_dataset, metric):
+        return update_model_if_needed(variant_dataset, metric)
+    
+    model_dropdown.change(
+        fn=on_model_change,
+        inputs=[model_dropdown, metric_dropdown],
+        outputs=[model_status]
+    )
+    
+    metric_dropdown.change(
+        fn=on_model_change,
+        inputs=[model_dropdown, metric_dropdown],
+        outputs=[model_status]
+    )
+
     submit_btn.click(
         fn=predict,
         inputs=[input_img, model_dropdown, metric_dropdown],