Update app.py

app.py

@@ -398,62 +398,84 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
 
 # ===== 分割功能 =====
 def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
-    """分割处理函数：输出对齐 overlay_instances 逻辑"""
+    """分割处理函数（与原始版本完全对齐，适配新版 run_seg 接口）"""
     if annot_value is None or len(annot_value) < 1:
+        print("❌ No annotation input")
         return None, "⚠️ 请先上传图像"
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
 
-    print(f"🖼️ Segmentation - Image: {img_path}")
-    
+    print(f"🖼️ Image path: {img_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)
+            box_array = [[xmin, ymin, xmax, ymax]]
             print(f"📦 Using box: {box_array}")
 
     try:
+        # 🔹 调用新版 run_seg 接口
         mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
-        if mask is None:
-            return None, "❌ 分割失败"
-        print(f"✅ Segmentation done - Mask shape: {mask.shape}")
+        print("📏 Mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
     except Exception as e:
-        print(f"❌ Segmentation error: {e}")
+        print(f"❌ Error during inference: {e}")
         return None, f"分割失败: {str(e)}"
 
     try:
-        # 读取原图（支持 JPG / PNG / TIF）
-        img = Image.open(img_path).convert("RGB")
-        img_rgb = np.array(img, dtype=np.float32)
-        if img_rgb.max() > 1.5:
-            img_rgb = img_rgb / 255.0
+        img = Image.open(img_path)
+        print("📷 Image mode:", img.mode, "size:", img.size)
     except Exception as e:
-        print(f"❌ Image load error: {e}")
-        return None, f"图像读取失败: {str(e)}"
+        print(f"❌ Failed to open image: {e}")
+        return None, f"图像打开失败: {str(e)}"
 
-    # === 实例 mask 可视化（与你提供的 overlay_instances 完全一致） ===
-    mask_np = np.array(mask).astype(np.int32)
-    overlay = img_rgb.copy()
-    cmap = cm.get_cmap("tab20", np.max(mask_np) + 1)
+    try:
+        img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
+        img_np = np.array(img_rgb, dtype=np.float32)
+        if img_np.max() > 1.5:
+            img_np = img_np / 255.0
+    except Exception as e:
+        print(f"❌ Error in image conversion/resizing: {e}")
+        return None, f"图像预处理失败: {str(e)}"
+
+    # 🔹 mask → 实例掩码
+    mask_np = np.array(mask)
+    inst_mask = mask_np.astype(np.int32)
+    unique_ids = np.unique(inst_mask)
+    num_instances = len(unique_ids[unique_ids != 0])
+    print(f"✅ Instance IDs found: {unique_ids}, Total instances: {num_instances}")
+
+    if num_instances == 0:
+        print("⚠️ No instance found, returning dummy red image")
+        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到细胞"
+
+    # ==== 🎨 Color Overlay (每个实例一个颜色) ====
+    overlay = img_np.copy()
     alpha = 0.5
+    cmap = cm.get_cmap("nipy_spectral", num_instances + 1)
 
-    for inst_id in np.unique(mask_np):
-        if inst_id == 0:  # 背景跳过
+    for inst_id in np.unique(inst_mask):
+        if inst_id == 0:
             continue
-        color = np.array(cmap(inst_id)[:3])  # RGB
-        overlay[mask_np == inst_id] = (1 - alpha) * overlay[mask_np == inst_id] + alpha * color
+        binary_mask = (inst_mask == inst_id).astype(np.uint8)
+        color = np.array(cmap(inst_id / (num_instances + 1))[:3])  # RGB only
+        overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color
+
+        # 🟡 绘制轮廓
+        contours = measure.find_contours(binary_mask, 0.5)
+        for contour in contours:
+            contour = contour.astype(np.int32)
+            overlay[contour[:, 0], contour[:, 1]] = [1.0, 1.0, 0.0]  # 黄色轮廓
 
     overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
-    overlay_img = Image.fromarray(overlay)
 
-    # === 返回统计结果 ===
-    num_instances = len(np.unique(mask_np)) - (1 if 0 in np.unique(mask_np) else 0)
-    result_text = f"✅ 检测到 {num_instances} 个细胞"
+    # 🔹 两种显示模式
+    if mode == "Instance Mask Only":
+        return Image.fromarray(colorize_mask(inst_mask, num_colors=512)), f"✅ 检测到 {num_instances} 个细胞"
 
-    return overlay_img, result_text
+    return Image.fromarray(overlay), f"✅ 检测到 {num_instances} 个细胞"
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):