Update app.py

app.py

@@ -397,8 +397,39 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
     return palette_arr[color_idx]
 
 # ===== 分割功能 =====
+# ===== 彩色 mask 可视化 =====
+def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
+    """将实例掩码转换为彩色图像，每个实例一个颜色"""
+    mask = mask.astype(np.int32)
+
+    def hsv_to_rgb(hh, ss, vv):
+        i = int(hh * 6.0)
+        f = hh * 6.0 - i
+        p = vv * (1.0 - ss)
+        q = vv * (1.0 - f * ss)
+        t = vv * (1.0 - (1.0 - f) * ss)
+        i = i % 6
+        if i == 0: r, g, b = vv, t, p
+        elif i == 1: r, g, b = q, vv, p
+        elif i == 2: r, g, b = p, vv, t
+        elif i == 3: r, g, b = p, q, vv
+        elif i == 4: r, g, b = t, p, vv
+        else: r, g, b = vv, p, q
+        return int(r*255), int(g*255), int(b*255)
+
+    palette = [(0, 0, 0)]
+    for k in range(1, num_colors):
+        hue = (k % num_colors) / float(num_colors)
+        palette.append(hsv_to_rgb(hue, 1.0, 0.95))
+
+    color_idx = mask % num_colors
+    palette_arr = np.array(palette, dtype=np.uint8)
+    return palette_arr[color_idx]
+
+
+# ===== 推理 + 实例彩色可视化 =====
 def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
-    """分割处理函数（与原始版本完全对齐，适配新版 run_seg 接口）"""
+    """分割推理 + 可视化，与原始逻辑完全对齐"""
     if annot_value is None or len(annot_value) < 1:
         print("❌ No annotation input")
         return None, "⚠️ 请先上传图像"
@@ -407,7 +438,6 @@ def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
     bboxes = annot_value[1] if len(annot_value) > 1 else []
 
     print(f"🖼️ Image path: {img_path}")
-
     box_array = None
     if use_box_choice == "Yes" and bboxes:
         box = parse_first_bbox(bboxes)
@@ -417,30 +447,25 @@ def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
             print(f"📦 Using box: {box_array}")
 
     try:
-        # 🔹 调用新版 run_seg 接口
+        # 调用分割模型
         mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
         print("📏 Mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
     except Exception as e:
         print(f"❌ Error during inference: {e}")
         return None, f"分割失败: {str(e)}"
 
+    # 读取原图
     try:
-        img = Image.open(img_path)
-        print("📷 Image mode:", img.mode, "size:", img.size)
-    except Exception as e:
-        print(f"❌ Failed to open image: {e}")
-        return None, f"图像打开失败: {str(e)}"
-
-    try:
-        img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
+        img = Image.open(img_path).convert("RGB")
+        img_rgb = img.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)}"
+        print(f"❌ Failed to open/convert image: {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)
@@ -451,31 +476,36 @@ def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
         print("⚠️ No instance found, returning dummy red image")
         return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到细胞"
 
-    # ==== 🎨 Color Overlay (每个实例一个颜色) ====
+    # ==== 🎨 Overlay (每个实例一个颜色 + 黄色轮廓) ====
     overlay = img_np.copy()
-    alpha = 0.5
-    cmap = cm.get_cmap("nipy_spectral", num_instances + 1)
+    alpha = 0.75  # 增强透明度，更清晰
+    cmap = cm.get_cmap("tab20", num_instances + 1)  # 亮色系更清楚
 
     for inst_id in np.unique(inst_mask):
         if inst_id == 0:
             continue
         binary_mask = (inst_mask == inst_id).astype(np.uint8)
-        color = np.array(cmap(inst_id / (num_instances + 1))[:3])  # RGB only
+        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]  # 黄色轮廓
+            for dy in [-1, 0, 1]:
+                for dx in [-1, 0, 1]:
+                    yy = np.clip(contour[:, 0] + dy, 0, overlay.shape[0] - 1).astype(np.int32)
+                    xx = np.clip(contour[:, 1] + dx, 0, overlay.shape[1] - 1).astype(np.int32)
+                    overlay[yy, xx] = [1.0, 1.0, 0.0]  # 粗边缘线
 
     overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
+    result_text = f"✅ 检测到 {num_instances} 个细胞"
 
-    # 🔹 两种显示模式
     if mode == "Instance Mask Only":
-        return Image.fromarray(colorize_mask(inst_mask, num_colors=512)), f"✅ 检测到 {num_instances} 个细胞"
+        # 输出纯彩色 mask
+        return Image.fromarray(colorize_mask(inst_mask, num_colors=512)), result_text
 
-    return Image.fromarray(overlay), f"✅ 检测到 {num_instances} 个细胞"
+    return Image.fromarray(overlay), result_text
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):