Update app.py

app.py

@@ -398,7 +398,7 @@ 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 逻辑"""
     if annot_value is None or len(annot_value) < 1:
         return None, "⚠️ 请先上传图像"
 
@@ -419,62 +419,41 @@ def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
         mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
         if mask is None:
             return None, "❌ 分割失败"
-        print(f"📏 Mask shape: {mask.shape}, dtype: {mask.dtype}, unique: {np.unique(mask)}")
+        print(f"✅ Segmentation done - Mask shape: {mask.shape}")
     except Exception as e:
         print(f"❌ Segmentation error: {e}")
         return None, f"分割失败: {str(e)}"
 
     try:
-        img = Image.open(img_path)
-        print(f"📷 Image mode: {img.mode}, size: {img.size}")
+        # 读取原图（支持 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
     except Exception as e:
-        print(f"❌ Failed to open image: {e}")
+        print(f"❌ Image load error: {e}")
         return None, f"图像读取失败: {str(e)}"
 
-    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_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:
-        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到细胞"
-
-    # ==== Color Overlay (每个实例一个颜色 + 黄色轮廓) ====
-    overlay = img_np.copy()
+    # === 实例 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)
     alpha = 0.5
-    cmap = cm.get_cmap("nipy_spectral", num_instances + 1)
 
-    for inst_id in unique_ids:
-        if inst_id == 0:
+    for inst_id in np.unique(mask_np):
+        if inst_id == 0:  # 背景跳过
             continue
-        binary_mask = (inst_mask == inst_id).astype(np.uint8)
-        color = np.array(cmap(inst_id / (num_instances + 1))[:3])
-        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]
+        color = np.array(cmap(inst_id)[:3])  # RGB
+        overlay[mask_np == inst_id] = (1 - alpha) * overlay[mask_np == inst_id] + alpha * color
 
     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)), result_text
-
-    return Image.fromarray(overlay), result_text
+    return overlay_img, result_text
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):