Update app.py

app.py

@@ -399,102 +399,102 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
 # ===== 分割功能 =====
 # ===== 彩色 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)
+    def hsv_to_rgb(h, s, v):
+        i = int(h * 6.0)
+        f = h * 6.0 - i
         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)
+        p = v * (1 - s)
+        q = v * (1 - f * s)
+        t = v * (1 - (1 - f) * s)
+        if i == 0: r, g, b = v, t, p
+        elif i == 1: r, g, b = q, v, p
+        elif i == 2: r, g, b = p, v, t
+        elif i == 3: r, g, b = p, q, v
+        elif i == 4: r, g, b = t, p, v
+        else: r, g, b = v, p, q
+        return int(r * 255), int(g * 255), int(b * 255)
+
+    palette = [(0, 0, 0)]  # 背景为黑色
+    for i in range(1, num_colors):
+        h = (i % num_colors) / float(num_colors)
+        palette.append(hsv_to_rgb(h, 1.0, 0.95))
 
-    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)
+    color_idx = mask % num_colors
     return palette_arr[color_idx]
 
+def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):
+    img = img.astype(np.float32)
+    if len(img.shape) == 2:
+        img = np.stack([img]*3, axis=-1)
+    if img.max() > 1.5:
+        img = img / 255.0
 
+    overlay = img.copy()
+    cmap = cm.get_cmap(cmap_name, np.max(mask) + 1)
+
+    for inst_id in np.unique(mask):
+        if inst_id == 0:
+            continue
+        color = np.array(cmap(inst_id)[:3])
+        overlay[mask == inst_id] = (1 - alpha) * overlay[mask == inst_id] + alpha * color
+
+    return overlay
 # ===== 推理 + 实例彩色可视化 =====
 def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
-    """分割推理 + 可视化（更深颜色、更清晰区分）"""
+    prepare_cuda()
     if annot_value is None or len(annot_value) < 1:
-        return None, "⚠️ 请先上传图像"
+        print("❌ No annotation input")
+        return None, "请上传图像"
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
 
+    print(f"🖼️ 图像路径: {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]]
+            print(f"📦 使用边界框: {box_array}")
 
+    # === Run model
     try:
-        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
-        if mask is None:
-            return None, "❌ 分割失败"
+        mask = run(MODEL, img_path, box=box_array, device=DEVICE)
+        print("📏 mask shape:", mask.shape, "unique ids:", np.unique(mask))
     except Exception as e:
-        return None, f"❌ 推理错误: {e}"
+        return None, f"❌ 推理失败: {str(e)}"
 
+    # === 读取原图
     try:
-        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)
+        img = Image.open(img_path).convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
+        img_np = np.array(img).astype(np.float32)
         if img_np.max() > 1.5:
-            img_np = img_np / 255.0
+            img_np /= 255.0
     except Exception as e:
-        return None, f"❌ 图像处理失败: {e}"
+        return None, f"❌ 图像读取失败: {str(e)}"
 
-    inst_mask = np.array(mask).astype(np.int32)
+    inst_mask = mask.astype(np.int32)
     unique_ids = np.unique(inst_mask)
     num_instances = len(unique_ids[unique_ids != 0])
+    print(f"✅ 实例数量: {num_instances}")
 
     if num_instances == 0:
-        return Image.new("RGB", mask.shape[::-1], (255, 200, 200)), "⚠️ 未检测到细胞"
+        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到实例"
 
-    # ===== 🎨 彩色 mask 生成 =====
-    color_mask = colorize_mask(inst_mask, num_colors=512).astype(np.float32) / 255.0
+    # === 可视化：Overlay 模式
+    if mode == "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, f"✅ 检测到 {num_instances} 个��胞"
 
-    # ===== 🌈 半透明叠加在原图上 =====
-    overlay = img_np.copy()
-    alpha = 0.6
-    blended = (1 - alpha) * overlay + alpha * color_mask
+    # === 可视化：纯彩色 mask
+    elif mode == "Instance Mask Only":
+        color_mask = colorize_mask(inst_mask, num_colors=512)
+        return Image.fromarray(color_mask), f"✅ 检测到 {num_instances} 个细胞"
 
-    # ===== 🔶 黄色轮廓线描边每个实例 =====
-    for inst_id in unique_ids:
-        if inst_id == 0:
-            continue
-        binary_mask = (inst_mask == inst_id).astype(np.uint8)
-        contours = measure.find_contours(binary_mask, 0.5)
-        for contour in contours:
-            contour = contour.astype(np.int32)
-            for dy in [-1, 0, 1]:
-                for dx in [-1, 0, 1]:
-                    yy = np.clip(contour[:, 0] + dy, 0, blended.shape[0] - 1).astype(np.int32)
-                    xx = np.clip(contour[:, 1] + dx, 0, blended.shape[1] - 1).astype(np.int32)
-                    blended[yy, xx] = [1.0, 1.0, 0.0]  # 黄色边
-
-    blended = np.clip(blended * 255.0, 0, 255).astype(np.uint8)
-    result_text = f"✅ 检测到 {num_instances} 个细胞"
-
-    if mode == "Instance Mask Only":
-        return Image.fromarray((color_mask * 255).astype(np.uint8)), result_text
-
-    return Image.fromarray(blended), result_text
+    return None, "❓ 无效显示模式"
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):