Update app.py

app.py

@@ -136,29 +136,11 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
     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):
-    """分割主函数 - 固定使用Overlay模式,同时返回原始mask"""
+    """分割主函数 - 每个实例不同颜色+轮廓"""
     if annot_value is None or len(annot_value) < 1:
+        print("❌ No annotation input")
         return None, None
 
     img_path = annot_value[0]
@@ -176,7 +158,7 @@ def segment_with_choice(use_box_choice, annot_value):
     # 运行分割模型
     try:
         mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
-        print("📏 mask shape:", mask.shape, "unique ids:", np.unique(mask))
+        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
     except Exception as e:
         print(f"❌ 推理失败: {str(e)}")
         return None, None
@@ -189,26 +171,55 @@ def segment_with_choice(use_box_choice, annot_value):
 
     # 读取原图
     try:
-        img = Image.open(img_path).convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
-        img_np = np.array(img).astype(np.float32)
+        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, None
+
+    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 /= 255.0
+            img_np = img_np / 255.0
     except Exception as e:
-        print(f"❌ 图像读取失败: {str(e)}")
+        print(f"❌ Error in image conversion/resizing: {e}")
         return None, None
 
-    inst_mask = mask.astype(np.int32)
+    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"✅ 实例数量: {num_instances}")
+    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)), 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, temp_mask_file.name
+    # ==== Color Overlay (每个实例一个颜色) ====
+    overlay = img_np.copy()
+    alpha = 0.5
+    cmap = cm.get_cmap("nipy_spectral", 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, ignore alpha
+        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)
+            # 确保坐标在范围内
+            valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
+            valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
+            overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]  # 黄色轮廓
+
+    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
+
+    return Image.fromarray(overlay), temp_mask_file.name
 
 # ===== 计数功能 =====
 def count_cells_handler(image_path):