From f0ac196afed1959c781ffec8b0eb8895d2c25c2c Mon Sep 17 00:00:00 2001
From: AbyssYuan0 <jiaxiangchang0@gmail.com>
Date: Mon, 8 Jan 2024 16:07:16 +0800
Subject: [PATCH] =?UTF-8?q?=E6=A0=B9=E6=8D=AE=E8=92=99=E7=89=88=E6=8A=A0?=
 =?UTF-8?q?=E5=9B=BE=E8=BF=94=E5=9B=9ERGBA=E7=9A=84=E5=9B=BE=E7=89=87?=
 =?UTF-8?q?=EF=BC=8C=E6=A0=B9=E6=8D=AE=E8=92=99=E7=89=88=E6=8A=A0=E5=9B=BE?=
 =?UTF-8?q?=E8=BF=94=E5=9B=9E=E6=89=A3=E4=B8=8BBox=E7=9A=84=E5=9B=BE?=
 =?UTF-8?q?=E7=89=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 __init__.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/__init__.py b/__init__.py
index cddf622..0462eed 100644
--- a/__init__.py
+++ b/__init__.py
@@ -615,6 +615,83 @@ class SegmentToMaskByPoint:
         return (mask0, mask1, mask2,)
 
 
+class CropImageByMask:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+            }
+        }
+
+    CATEGORY = "badger"
+
+    RETURN_TYPES = ("IMAGE", "INT", "INT",)
+    RETURN_NAMES = ("cropped_img", "X", "Y",)
+    FUNCTION = "crop_image_by_mask"
+
+    def crop_image_by_mask(self, image, mask):
+        # Ensure the mask is binary
+        mask = (mask > 0.5).float()
+
+        # Find the bounding box of the mask
+        if mask.sum() == 0:
+            raise ValueError("The mask is empty, cannot determine bounding box for cropping.")
+
+        # Find indices where the mask is nonzero
+        nonzero_indices = torch.nonzero(mask.squeeze(0), as_tuple=True)
+        topmost = torch.min(nonzero_indices[0])
+        leftmost = torch.min(nonzero_indices[1])
+        bottommost = torch.max(nonzero_indices[0])
+        rightmost = torch.max(nonzero_indices[1])
+
+        # Crop the image using the bounding box
+        cropped_image = image[:, topmost:bottommost + 1, leftmost:rightmost + 1]
+
+        # Return the cropped image and the top-left coordinates of the bounding box
+        X = int(leftmost)
+        Y = int(topmost)
+        return (cropped_image, X, Y,)
+
+
+class ApplyMaskToImage:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+            }
+        }
+
+    CATEGORY = "badger"
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("rgba_image",)
+    FUNCTION = "apply_mask_to_image"
+
+    def apply_mask_to_image(self, image, mask):
+        image = tensorToImg(image)
+        mask = tensorToImg(mask)
+        mask = mask.convert("L")
+
+        # 将图片转换为RGBA，以便添加透明度通道
+        image = image.convert("RGBA")
+
+        # 分离图片的通道
+        r, g, b, a = image.split()
+
+        # 将蒙版应用为alpha通道
+        new_a = Image.composite(a, Image.new('L', mask.size, 0), mask)
+
+        # 合并图像通道和新的alpha通道
+        result_image = Image.merge('RGBA', (r, g, b, new_a))
+        return (imgToTensor(result_image),)
+
+
+
 NODE_CLASS_MAPPINGS = {
     "ImageOverlap-badger": ImageOverlap,
     "FloatToInt-badger": FloatToInt,
@@ -630,6 +707,8 @@ NODE_CLASS_MAPPINGS = {
     "mkdir-badger": mkdir,
     "findCenterOfMask-badger": findCenterOfMask,
     "SegmentToMaskByPoint-badger": SegmentToMaskByPoint,
+    "CropImageByMask-badger": CropImageByMask,
+    "ApplyMaskToImage-badger": ApplyMaskToImage,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {