2025-07-31 21:55 🚀

app.py

@@ -25,22 +25,46 @@ 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",
-    "ZIP-S @ ShanghaiTech A", "ZIP-S @ ShanghaiTech B", "ZIP-S @ UCF-QNRF",
-    "ZIP-T @ ShanghaiTech A", "ZIP-T @ ShanghaiTech B", "ZIP-T @ UCF-QNRF",
-    "ZIP-N @ ShanghaiTech A", "ZIP-N @ ShanghaiTech B", "ZIP-N @ UCF-QNRF",
-    "ZIP-P @ ShanghaiTech A", "ZIP-P @ ShanghaiTech B", "ZIP-P @ UCF-QNRF"
+    "ZIP-B @ ShanghaiTech A @ MAE", "ZIP-B @ ShanghaiTech A @ RMSE", "ZIP-B @ ShanghaiTech A @ NAE",
+    "ZIP-B @ ShanghaiTech B @ MAE", "ZIP-B @ ShanghaiTech B @ RMSE", "ZIP-B @ ShanghaiTech B @ NAE",
+    "ZIP-B @ UCF-QNRF @ MAE", "ZIP-B @ UCF-QNRF @ RMSE", "ZIP-B @ UCF-QNRF @ NAE",
+    "ZIP-B @ NWPU-Crowd @ MAE", "ZIP-B @ NWPU-Crowd @ RMSE", "ZIP-B @ NWPU-Crowd @ NAE",
+    "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+    "ZIP-S @ ShanghaiTech A @ MAE", "ZIP-S @ ShanghaiTech A @ RMSE", "ZIP-S @ ShanghaiTech A @ NAE",
+    "ZIP-S @ ShanghaiTech B @ MAE", "ZIP-S @ ShanghaiTech B @ RMSE", "ZIP-S @ ShanghaiTech B @ NAE",
+    "ZIP-S @ UCF-QNRF @ MAE", "ZIP-S @ UCF-QNRF @ RMSE", "ZIP-S @ UCF-QNRF @ NAE",
+    "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+    "ZIP-T @ ShanghaiTech A @ MAE", "ZIP-T @ ShanghaiTech A @ RMSE", "ZIP-T @ ShanghaiTech A @ NAE",
+    "ZIP-T @ ShanghaiTech B @ MAE", "ZIP-T @ ShanghaiTech B @ RMSE", "ZIP-T @ ShanghaiTech B @ NAE",
+    "ZIP-T @ UCF-QNRF @ MAE", "ZIP-T @ UCF-QNRF @ RMSE", "ZIP-T @ UCF-QNRF @ NAE",
+    "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+    "ZIP-N @ ShanghaiTech A @ MAE", "ZIP-N @ ShanghaiTech A @ RMSE", "ZIP-N @ ShanghaiTech A @ NAE",
+    "ZIP-N @ ShanghaiTech B @ MAE", "ZIP-N @ ShanghaiTech B @ RMSE", "ZIP-N @ ShanghaiTech B @ NAE",
+    "ZIP-N @ UCF-QNRF @ MAE", "ZIP-N @ UCF-QNRF @ RMSE", "ZIP-N @ UCF-QNRF @ NAE",
+    "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+    "ZIP-P @ ShanghaiTech A @ MAE", "ZIP-P @ ShanghaiTech A @ RMSE", "ZIP-P @ ShanghaiTech A @ NAE",
+    "ZIP-P @ ShanghaiTech B @ MAE", "ZIP-P @ ShanghaiTech B @ RMSE", "ZIP-P @ ShanghaiTech B @ NAE",
+    "ZIP-P @ UCF-QNRF @ MAE", "ZIP-P @ UCF-QNRF @ RMSE", "ZIP-P @ UCF-QNRF @ NAE",
 ]
 
 # -----------------------------
 # Model management functions
 # -----------------------------
-def update_model_if_needed(variant_dataset: str, metric: str):
+def update_model_if_needed(variant_dataset_metric: str):
     """
     Load a new model only if the configuration has changed.
     """
     global loaded_model, current_model_config
-    variant, dataset = variant_dataset.split(" @ ")
+    
+    # 如果是分割线，则跳过
+    if "━━━━━━" in variant_dataset_metric:
+        return "Please select a valid model configuration"
+    
+    parts = variant_dataset_metric.split(" @ ")
+    if len(parts) != 3:
+        return "Invalid model configuration format"
+        
+    variant, dataset, metric = parts[0], parts[1], parts[2].lower()
 
     if dataset == "ShanghaiTech A":
         dataset_name = "sha"
@@ -50,6 +74,8 @@ def update_model_if_needed(variant_dataset: str, metric: str):
         dataset_name = "qnrf"
     elif dataset == "NWPU-Crowd":
         dataset_name = "nwpu"
+    else:
+        return f"Unknown dataset: {dataset}"
 
     if (loaded_model is None or 
         current_model_config["variant"] != variant or 
@@ -271,17 +297,25 @@ def _sliding_window_predict(
 # Inference function
 # -----------------------------
 @spaces.GPU(duration=120)
-def predict(image: Image.Image, variant_dataset: str, metric: str):
+def predict(image: Image.Image, variant_dataset_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, current_model_config
     
+    # 如果选择的是分割线，返回错误信息
+    if "━━━━━━" in variant_dataset_metric:
+        return image, None, None, "Please select a valid model configuration", None, None, None
+    
     # 确保模型正确加载
-    update_model_if_needed(variant_dataset, metric)
+    update_model_if_needed(variant_dataset_metric)
     
-    variant, dataset = variant_dataset.split(" @ ")
+    parts = variant_dataset_metric.split(" @ ")
+    if len(parts) != 3:
+        return image, None, None, "Invalid model configuration format", None, None, None
+        
+    variant, dataset, metric = parts[0], parts[1], parts[2].lower()
 
     if dataset == "ShanghaiTech A":
         dataset_name = "sha"
@@ -291,6 +325,8 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
         dataset_name = "qnrf"
     elif dataset == "NWPU-Crowd":
         dataset_name = "nwpu"
+    else:
+        return image, None, None, f"Unknown dataset: {dataset}", None, None, None
 
     if not hasattr(loaded_model, "input_size"):
         if dataset_name == "sha":
@@ -424,7 +460,57 @@ def predict(image: Image.Image, variant_dataset: str, metric: str):
 # -----------------------------
 # Build Gradio Interface using Blocks for a two-column layout
 # -----------------------------
-with gr.Blocks() as demo:
+css = """
+/* 分割线样式 - 灰色不可选择 */
+.dropdown select option[value*="━━━━━━"] {
+    color: #999 !important;
+    background-color: #f0f0f0 !important;
+    font-style: italic !important;
+    text-align: center !important;
+    pointer-events: none !important;
+    cursor: not-allowed !important;
+    border: none !important;
+}
+
+/* Gradio下拉菜单中的分割线样式 */
+.gr-dropdown .choices__item[data-value*="━━━━━━"] {
+    color: #999 !important;
+    background-color: #f0f0f0 !important;
+    font-style: italic !important;
+    text-align: center !important;
+    pointer-events: none !important;
+    cursor: not-allowed !important;
+    user-select: none !important;
+    opacity: 0.6 !important;
+}
+
+/* 悬停时保持灰色 */
+.gr-dropdown .choices__item[data-value*="━━━━━━"]:hover {
+    background-color: #f0f0f0 !important;
+    color: #999 !important;
+    cursor: not-allowed !important;
+}
+
+/* 通用的分割线样式 */
+option:disabled {
+    color: #999 !important;
+    background-color: #f0f0f0 !important;
+    font-style: italic !important;
+}
+
+/* 为包含分割线字符的选项添加样式 */
+option[value*="━━━━━━"], 
+select option[value*="━━━━━━"] {
+    color: #999 !important;
+    background-color: #f0f0f0 !important;
+    cursor: not-allowed !important;
+    pointer-events: none !important;
+    text-align: center !important;
+    opacity: 0.6 !important;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
     gr.Markdown("# Crowd Counting by ZIP")
     gr.Markdown("Upload an image or select an example below to see the predicted crowd density map and total count.")
 
@@ -433,17 +519,9 @@ with gr.Blocks() as demo:
             # Dropdown for model variant
             model_dropdown = gr.Dropdown(
                 choices=pretrained_models,
-                value="ZIP-B @ NWPU-Crowd",
+                value="ZIP-B @ NWPU-Crowd @ MAE",
                 label="Select a pretrained model"
             )
-
-            # Dropdown for metric, always the same choices
-            metric_dropdown = gr.Dropdown(
-                choices=["mae", "rmse", "nae"],
-                value="mae",
-                label="Select Best Metric"
-            )
-
             model_status = gr.Textbox(
                 label="Model Status",
                 value="No model loaded",
@@ -463,31 +541,28 @@ 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)
+    # 当模型变化时，自动更新模型
+    def on_model_change(variant_dataset_metric):
+        # 如果选择的是分割线，保持当前选择不变
+        if "━━━━━━" in variant_dataset_metric:
+            return "Please select a valid model configuration"
+        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],
+        inputs=[model_dropdown],
         outputs=[model_status]
     )
 
     # 页面加载时自动加载默认模型
     demo.load(
-        fn=lambda: update_model_if_needed("ZIP-B @ NWPU-Crowd", "mae"),
+        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],
+        inputs=[input_img, model_dropdown],
         outputs=[input_img, output_den_map, output_lambda_map, output_text, output_structural_zero_map, output_sampling_zero_map, output_complete_zero_map]
     )