Update app.py

app.py

@@ -157,9 +157,9 @@ def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):
 
 # ===== 分割功能 =====
 def segment_with_choice(use_box_choice, annot_value):
-    """分割主函数 - 固定使用Overlay模式"""
+    """分割主函数 - 固定使用Overlay模式,同时返回原始mask"""
     if annot_value is None or len(annot_value) < 1:
-        return None
+        return None, None
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
@@ -179,7 +179,13 @@ def segment_with_choice(use_box_choice, annot_value):
         print("📏 mask shape:", mask.shape, "unique ids:", np.unique(mask))
     except Exception as e:
         print(f"❌ 推理失败: {str(e)}")
-        return None
+        return None, None
+
+    # 保存原始mask为TIF文件
+    temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tif")
+    mask_img = Image.fromarray(mask.astype(np.uint16))
+    mask_img.save(temp_mask_file.name)
+    print(f"💾 原始mask保存到: {temp_mask_file.name}")
 
     # 读取原图
     try:
@@ -189,7 +195,7 @@ def segment_with_choice(use_box_choice, annot_value):
             img_np /= 255.0
     except Exception as e:
         print(f"❌ 图像读取失败: {str(e)}")
-        return None
+        return None, None
 
     inst_mask = mask.astype(np.int32)
     unique_ids = np.unique(inst_mask)
@@ -197,16 +203,16 @@ def segment_with_choice(use_box_choice, annot_value):
     print(f"✅ 实例数量: {num_instances}")
 
     if num_instances == 0:
-        return Image.new("RGB", mask.shape[::-1], (255, 0, 0))
+        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None
 
     # 使用Overlay模式可视化
     overlay = overlay_instances(img_np, inst_mask, alpha=0.5, cmap_name="tab20")
     overlay_img = Image.fromarray((overlay * 255).astype(np.uint8))
-    return overlay_img
+    return overlay_img, temp_mask_file.name
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):
-    """计数处理函数"""
+    """计数处理函数 - 只返回密度图"""
     if image_path is None:
         return None, "⚠️ 请先上传图像"
     
@@ -225,12 +231,34 @@ def count_cells_handler(image_path):
             return None, f"❌ 计数失败: {result['error']}"
         
         count = result['count']
-        viz_path = result['visualized_path']
+        
+        # 只提取密度图部分(假设visualized_path是拼接图,我们只要右半部分)
+        viz_path = result.get('visualized_path')
+        
+        # 如果有density_map_path,直接使用
+        if 'density_map_path' in result:
+            density_path = result['density_map_path']
+        elif viz_path and os.path.exists(viz_path):
+            # 如果是拼接图,提取右半部分(密度图)
+            try:
+                viz_img = Image.open(viz_path)
+                w, h = viz_img.size
+                # 取右半部分
+                density_img = viz_img.crop((w//2, 0, w, h))
+                # 保存为新文件
+                temp_density = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+                density_img.save(temp_density.name)
+                density_path = temp_density.name
+            except:
+                density_path = viz_path
+        else:
+            density_path = viz_path
+        
         result_text = f"✅ 检测到 {count:.1f} 个细胞"
         
         print(f"✅ Counting done - Count: {count:.1f}")
         
-        return viz_path, result_text
+        return density_path, result_text
         
     except Exception as e:
         print(f"❌ Counting error: {e}")
@@ -369,6 +397,12 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         height=400
                     )
                     
+                    # 下载原始预测结果
+                    download_mask_btn = gr.File(
+                        label="📥 下载原始预测 (.tif 格式)",
+                        visible=True
+                    )
+                    
                     # 满意度评分
                     score_slider = gr.Slider(
                         minimum=1,
@@ -398,7 +432,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
             run_seg_btn.click(
                 fn=segment_with_choice,
                 inputs=[use_box_radio, annotator],
-                outputs=seg_output
+                outputs=[seg_output, download_mask_btn]
             )
             
             # 初始化Gallery显示
@@ -486,7 +520,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                 
                 with gr.Column(scale=2):
                     count_output = gr.Image(
-                        label="📸 计数结果",
+                        label="📸 密度图",
                         type="filepath",
                         height=500
                     )