2025-07-31 19:44 🐛

Fixed bugs in app.py

app.py

@@ -23,13 +23,13 @@ EPS = 1e-8
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 loaded_model = None
 
-pretrained_datasets = {
-    "ZIP-B": ["ShanghaiTech A", "ShanghaiTech B", "UCF-QNRF", "NWPU-Crowd"],
-    "ZIP-S": ["ShanghaiTech A", "ShanghaiTech B", "UCF-QNRF"],
-    "ZIP-T": ["ShanghaiTech A", "ShanghaiTech B", "UCF-QNRF"],
-    "ZIP-N": ["ShanghaiTech A", "ShanghaiTech B", "UCF-QNRF"],
-    "ZIP-P": ["ShanghaiTech A", "ShanghaiTech B", "UCF-QNRF"],
-}
+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"
+]
 
 # -----------------------------
 # Define the model architecture
@@ -386,17 +386,10 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             # Dropdown for model variant
-            variant_dropdown = gr.Dropdown(
-                choices=list(pretrained_datasets.keys()),
-                value="ZIP-B",
-                label="Select Model Variant"
-            )
-
-            # Dropdown for pretrained dataset, dynamically updated based on variant
-            dataset_dropdown = gr.Dropdown(
-                choices=pretrained_datasets["ZIP-B"],
-                value=pretrained_datasets["ZIP-B"][0],
-                label="Select Pretrained Dataset"
+            model_dropdown = gr.Dropdown(
+                choices=pretrained_models,
+                value="ZIP-B @ NWPU-Crowd",
+                label="Select a pretrained model"
             )
 
             # Dropdown for metric, always the same choices
@@ -406,19 +399,6 @@ with gr.Blocks() as demo:
                 label="Select Best Metric"
             )
 
-            # Update dataset choices when variant changes
-            def update_dataset(variant):
-                choices = pretrained_datasets[variant]
-                return gr.Dropdown.update(
-                    choices=choices,
-                    value=choices[0]
-                )
-
-            variant_dropdown.change(
-                fn=update_dataset,
-                inputs=variant_dropdown,
-                outputs=dataset_dropdown
-            )
             input_img = gr.Image(label="Input Image", sources=["upload", "clipboard"], type="pil")
             submit_btn = gr.Button("Predict")
     
@@ -431,9 +411,10 @@ with gr.Blocks() as demo:
             
             output_text = gr.Textbox(label="Total Count")
 
+    variant, dataset = model_dropdown.value.split(" @ ")
     submit_btn.click(
         fn=predict,
-        inputs=[input_img, variant_dropdown, dataset_dropdown, metric_dropdown],
+        inputs=[input_img, variant, dataset, metric_dropdown],
         outputs=[input_img, output_structural_zero_map, output_sampling_zero_map, output_complete_zero_map, output_lambda_map, output_den_map, output_text]
     )