Update seg_script.py

seg_script.py

@@ -39,16 +39,13 @@ def process_image(image_path, output_dir):
     drawed = img.copy()
     im_h, im_w = img.shape[:2]
 
+    # 创建黑色背景的 mask 图像
+    mask_image = np.zeros((im_h, im_w), dtype=np.uint8)
+
     # 如果没有检测到对象，直接返回原图
     if instances.bboxes is None:
         return
 
-    # 保存绘制后的图像（只保存一次）
-    base_name = Path(image_path).stem
-    output_path = Path(output_dir)
-    output_path.mkdir(parents=True, exist_ok=True)
-    Image.fromarray(drawed).save(output_path / f"{base_name}_drawed.png")
-
     # 处理每个实例
     for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
         color = get_color(ii)
@@ -69,6 +66,9 @@ def process_image(image_path, output_dir):
 
         drawed = drawed.astype(np.uint8)
 
+        # 将当前实例的 mask 绘制到黑色背景上（白色表示 mask）
+        mask_image[mask > 0] = 255
+
         # 裁剪图像
         x1, y1, x2, y2 = int(xywh[0]), int(xywh[1]), int(xywh[0] + xywh[2]), int(xywh[1] + xywh[3])
         cropped_img = img[y1:y2, x1:x2]
@@ -78,10 +78,23 @@ def process_image(image_path, output_dir):
         alpha_channel = (cropped_mask * 255).astype(np.uint8)
         rgba_image = np.dstack((cropped_img, alpha_channel))
 
-        # 保存裁剪后的图像和分割后的图像（文件名包含实例下标）
+        # 保存裁剪后的图像和分割后的图像
+        base_name = Path(image_path).stem
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        # 保存裁剪后的图像
         Image.fromarray(cropped_img).save(output_path / f"{base_name}_cropped_{ii}.png")
+
+        # 保存分割后的图像（RGBA）
         Image.fromarray(rgba_image, 'RGBA').save(output_path / f"{base_name}_segmented_{ii}.png")
 
+    # 保存绘制后的图像
+    Image.fromarray(drawed).save(output_path / f"{base_name}_drawed.png")
+
+    # 保存 mask 图像
+    Image.fromarray(mask_image).save(output_path / f"{base_name}_mask.png")
+
 def main():
     parser = argparse.ArgumentParser(description="Anime Instance Segmentation")
     parser.add_argument("input_path", type=str, help="Path to the input image or folder")