2025-07-31 21:38 🐛

app.py

@@ -277,6 +277,10 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
     compute the total crowd count, and prepare the density map for display.
     """
     global loaded_model, current_model_config
+    
+    # 确保模型正确加载
+    update_model_if_needed(variant_dataset, metric)
+    
     variant, dataset = variant_dataset.split(" @ ")
 
     if dataset == "ShanghaiTech A":
@@ -288,17 +292,6 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
     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}
-    else:
-        print(f"Using cached model: {variant} @ {dataset} with {metric} metric")
-
     if not hasattr(loaded_model, "input_size"):
         if dataset_name == "sha":
             loaded_model.input_size = 224
@@ -470,6 +463,7 @@ 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)
     
@@ -485,6 +479,12 @@ with gr.Blocks() as demo:
         outputs=[model_status]
     )
 
+    # 页面加载时自动加载默认模型
+    demo.load(
+        fn=lambda: update_model_if_needed("ZIP-B @ NWPU-Crowd", "mae"),
+        outputs=[model_status]
+    )
+
     submit_btn.click(
         fn=predict,
         inputs=[input_img, model_dropdown, metric_dropdown],