From 3b074e256c92152f350cb044cb9529b41fd15a5b Mon Sep 17 00:00:00 2001
From: yada <trojblue@gmail.com>
Date: Fri, 31 Mar 2023 03:42:09 -0400
Subject: [PATCH] initial commit: add image layering

---
 __init__.py       |  6 +++
 debug.py          | 37 ++++++++++++++++++
 image_layering.py | 98 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 141 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 debug.py
 create mode 100644 image_layering.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..8019d87
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,6 @@
+import custom_nodes.comfy_nodes_trojblue.image_layering as image_layering
+
+
+NODE_CLASS_MAPPINGS = {
+    "layering": image_layering.Layering,                      # Layering
+}
diff --git a/debug.py b/debug.py
new file mode 100644
index 0000000..f2b4741
--- /dev/null
+++ b/debug.py
@@ -0,0 +1,37 @@
+from PIL import Image
+import numpy as np
+
+
+
+def alpha_cutout(img, threshold=100, dist=5):
+    arr = np.array(np.asarray(img))  # 获取图像数据，使用了numpy
+    r, g, b, a = np.rollaxis(arr, axis=-1)
+
+    mask = ((r > threshold)
+            & (g > threshold)
+            & (b > threshold)
+            & (np.abs(r - g) < dist)  # 将接近白色背景的也替换掉
+            & (np.abs(r - b) < dist)
+            & (np.abs(g - b) < dist)
+            )
+    arr[mask, 3] = 0
+
+    img = Image.fromarray(arr, mode='RGBA')  # 转换为图像格式
+    return img
+
+def cutout_csdn():
+
+    threshold = 100
+    dist = 5
+    img = Image.open("img.png").convert('RGBA')  # 增加Alpha通道
+
+    img = alpha_cutout(img, threshold, dist)
+
+
+    img.show()
+
+
+if __name__ == '__main__':
+    cutout_csdn()
+    # image = "img.png"
+    # do_cutout(image)
diff --git a/image_layering.py b/image_layering.py
new file mode 100644
index 0000000..7b1e276
--- /dev/null
+++ b/image_layering.py
@@ -0,0 +1,98 @@
+import torch
+import numpy as np
+from PIL import Image
+import subprocess
+import sys
+try:
+    import blend_modes
+except ModuleNotFoundError:
+    # install pixelsort in current venv
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "blend-modes"])
+    import blend_modes
+
+import torch
+import numpy as np
+from PIL import Image
+
+class Layering:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_image": ("IMAGE",),
+                "add_layer1": ("IMAGE",)},
+            "optional": {
+                "add_layer2": ("IMAGE", {"default": None}),
+                "add_layer3": ("IMAGE", {"default": None}),
+                "key_color": ("TUPLE", {"default": (255, 255, 255)}),
+                # "alpha1": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                # "alpha2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                # "alpha3": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_blend"
+    CATEGORY = "trojblue_folder"
+
+    def tensor_to_pil(self, img):
+        if img is not None:
+            i = 255. * img.cpu().numpy().squeeze()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        return img
+
+    def alpha_cutout(self, img, threshold=80, dist=10):
+        arr = np.array(np.asarray(img))  # 获取图像数据，使用了numpy
+        r, g, b, a = np.rollaxis(arr, axis=-1)
+
+        mask = ((r > threshold)
+                & (g > threshold)
+                & (b > threshold)
+                & (np.abs(r - g) < dist)  # 将接近白色背景的也替换掉
+                & (np.abs(r - b) < dist)
+                & (np.abs(g - b) < dist)
+                )
+        arr[mask, 3] = 0
+
+        img = Image.fromarray(arr, mode='RGBA')  # 转换为图像格式
+        return img
+
+
+    def create_transparent_image(self, image, key_color, alpha):
+        transparent_image = Image.new('RGBA', image.size, (0, 0, 0, 0))
+        for x in range(image.width):
+            for y in range(image.height):
+                pixel = image.getpixel((x, y))
+                if pixel != key_color:
+                    transparent_image.putpixel((x, y), (*pixel[:3], int(255 * alpha)))
+        return transparent_image
+
+
+    def apply_blend(self, base_image, add_layer1, alpha1, add_layer2=None, alpha2=1.0, add_layer3=None, alpha3=1.0, key_color=(255, 255, 255)):
+        base_image = self.tensor_to_pil(base_image[0]).convert('RGBA')
+        add_layers = [(add_layer1, alpha1), (add_layer2, alpha2), (add_layer3, alpha3)]
+        add_layers = [(self.tensor_to_pil(layer[0]).convert('RGBA'), alpha) for layer, alpha in add_layers if layer is not None]
+
+        for image, alpha in add_layers:
+            image = image.resize(base_image.size, Image.ANTIALIAS)
+            transparent_image = self.alpha_cutout(image)
+            base_image = Image.alpha_composite(base_image, transparent_image)
+
+        base_image = base_image.convert('RGB')
+
+        # convert to tensor
+        out_image = np.array(base_image).astype(np.float32) / 255.0
+        out_image = torch.from_numpy(out_image).unsqueeze(0)
+
+        return (out_image,)
+
+NODE_CLASS_MAPPINGS = {
+    "Layering": Layering,
+}
+
+
+