2025-07-31 21:23 🐛

Fixed a bug in app.py

app.py

@@ -36,7 +36,7 @@ pretrained_models = [
 # -----------------------------
 def load_model(variant: str, dataset: str = "ShanghaiTech B", metric: str = "mae"):
     """ Load the model weights from the Hugging Face Hub."""
-    global loaded_model
+    # global loaded_model
     # Build model
 
     model_info_path = hf_hub_download(
@@ -241,31 +241,30 @@ 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
     variant, dataset = variant_dataset.split(" @ ")
 
-    if loaded_model is None:
+    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 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:
         load_model(variant=variant, dataset=dataset_name, metric=metric)
 
-        if not hasattr(loaded_model, "input_size"):
-            if dataset_name == "sha":
-                loaded_model.input_size = 224
-            elif dataset_name == "shb":
-                loaded_model.input_size = 448
-            elif dataset_name == "qnrf":
-                loaded_model.input_size = 672
-            elif dataset_name == "nwpu":
-                loaded_model.input_size = 672
+    if not hasattr(loaded_model, "input_size"):
+        if dataset_name == "sha":
+            loaded_model.input_size = 224
+        elif dataset_name == "shb":
+            loaded_model.input_size = 448
+        elif dataset_name == "qnrf":
+            loaded_model.input_size = 672
+        elif dataset_name == "nwpu":
+            loaded_model.input_size = 672
 
     loaded_model.to(device)