Update app.py

app.py

@@ -222,18 +222,30 @@ def segment_with_choice(use_box_choice, annot_value):
     return Image.fromarray(overlay), temp_mask_file.name
 
 # ===== 计数功能 =====
-def count_cells_handler(image_path):
-    """计数处理函数 - 只返回密度图"""
-    if image_path is None:
+def count_cells_handler(use_box_choice, annot_value):
+    """计数处理函数 - 支持边界框,只返回密度图"""
+    if annot_value is None or len(annot_value) < 1:
         return None, "⚠️ 请先上传图像"
     
+    image_path = annot_value[0]
+    bboxes = annot_value[1] if len(annot_value) > 1 else []
+    
+    print(f"🖼️ 图像路径: {image_path}")
+    box_array = None
+    if use_box_choice == "Yes" and bboxes:
+        box = parse_first_bbox(bboxes)
+        if box:
+            xmin, ymin, xmax, ymax = map(int, box)
+            box_array = [[xmin, ymin, xmax, ymax]]
+            print(f"📦 使用边界框: {box_array}")
+    
     try:
         print(f"🔢 Counting - Image: {image_path}")
         
         result = run_count(
             COUNT_MODEL,
             image_path,
-            box=None,
+            box=box_array,
             device=COUNT_DEVICE,
             visualize=True
         )
@@ -514,9 +526,9 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
             
             with gr.Row():
                 with gr.Column(scale=1):
-                    count_input = gr.Image(
-                        label="🖼️ 上传图像",
-                        type="filepath"
+                    count_annotator = BBoxAnnotator(
+                        label="🖼️ 上传图像 (可选标注边界框)",
+                        categories=["cell"]
                     )
                     
                     # 示例图片Gallery (与Segmentation相同)
@@ -533,14 +545,21 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         type="filepath"
                     )
                     
+                    with gr.Row():
+                        count_use_box_radio = gr.Radio(
+                            choices=["Yes", "No"],
+                            value="No",
+                            label="🔲 使用边界框?"
+                        )
+                    
                     count_btn = gr.Button("▶️ 运行计数", variant="primary", size="lg")
                     
                     gr.Markdown(
                         """
                         **使用说明:**
                         1. 上传图像或从Gallery选择示例
-                        2. 点击 "运行计数"
-                        3. 查看密度图和计数结果
+                        2. (可选) 标注边界框并选择 "Yes"
+                        3. 点击 "运行计数"
                         """
                     )
                 
@@ -558,7 +577,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
             # 绑定事件
             count_btn.click(
                 fn=count_cells_handler,
-                inputs=count_input,
+                inputs=[count_use_box_radio, count_annotator],
                 outputs=[count_output, count_status]
             )
             
@@ -591,7 +610,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                 outputs=count_example_gallery
             )
             
-            # 绑定事件: 点击Gallery加载到count_input
+            # 绑定事件: 点击Gallery加载到count_annotator
             def load_from_count_gallery(evt: gr.SelectData, all_imgs):
                 if evt.index is not None and evt.index < len(all_imgs):
                     return all_imgs[evt.index]
@@ -600,7 +619,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
             count_example_gallery.select(
                 fn=load_from_count_gallery,
                 inputs=count_user_examples,
-                outputs=count_input
+                outputs=count_annotator
             )
         
         # ===== Tab 3: Tracking =====