From 05144865d753e71ad817b23d99812adf1cec5c52 Mon Sep 17 00:00:00 2001
From: yada <trojblue@gmail.com>
Date: Fri, 31 Mar 2023 06:10:39 -0400
Subject: [PATCH] add color correction & image layering node

---
 README.md           | 18 +++++++++++++
 __init__.py         |  3 +++
 color_correction.py | 61 +++++++++++++++++++++++++++++++++++++++++++++
 image_layering.py   |  4 +--
 4 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 color_correction.py

diff --git a/README.md b/README.md
index e482c6c..b10ea20 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,20 @@
+
 # comfy-custom-nodes
+
 custom node modules for ComfyUI
+
+## Installation
+
+```bash
+cd python_embeded
+./python.exe -m pip install -r opencv-python scikit-image blendmodes
+```
+
+## Nodes
+
+image_layering:
+- Adds 1-3 layers of image on top of one another; will remove white background and use it as transparent layers
+
+color_correction:
+- Adjusts the color of the target image according to another image; ported from stable diffusion WebUI
+
diff --git a/__init__.py b/__init__.py
index 8019d87..e925529 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,6 +1,9 @@
 import custom_nodes.comfy_nodes_trojblue.image_layering as image_layering
+import custom_nodes.comfy_nodes_trojblue.color_correction as color_correction
 
 
 NODE_CLASS_MAPPINGS = {
     "layering": image_layering.Layering,                      # Layering
+    "color_correction": color_correction.ColorCorrectionNode,  # ColorCorrectionNode
+
 }
diff --git a/color_correction.py b/color_correction.py
new file mode 100644
index 0000000..d0ab8b0
--- /dev/null
+++ b/color_correction.py
@@ -0,0 +1,61 @@
+import torch
+from PIL import Image
+import numpy as np
+import cv2
+from skimage import exposure
+from blendmodes.blend import blendLayers, BlendType
+
+class ColorCorrectionNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "original_image": ("IMAGE",),
+                "target_image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "color_correct"
+    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 apply_color_correction(self, correction, original_image):
+
+        # https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/22bcc7be428c94e9408f589966c2040187245d81/modules/processing.py#L44
+
+        correction_target = cv2.cvtColor(np.asarray(correction.copy()), cv2.COLOR_RGB2LAB)
+
+        image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
+            cv2.cvtColor(
+                np.asarray(original_image),
+                cv2.COLOR_RGB2LAB
+            ),
+            correction_target,
+            channel_axis=2
+        ), cv2.COLOR_LAB2RGB).astype("uint8"))
+
+        image = blendLayers(image, original_image, BlendType.LUMINOSITY)
+        return image
+
+    def color_correct(self, original_image, target_image):
+        original_image = self.tensor_to_pil(original_image)
+        target_image = self.tensor_to_pil(target_image)
+
+        corrected_image = self.apply_color_correction(target_image, original_image)
+
+        # convert to tensor
+        corrected_image = corrected_image.convert('RGB')
+        out_image = np.array(corrected_image).astype(np.float32) / 255.0
+        out_image = torch.from_numpy(out_image).unsqueeze(0)
+
+        return (out_image,)
+
+NODE_CLASS_MAPPINGS = {
+    "ColorCorrectionNode": ColorCorrectionNode
+}
diff --git a/image_layering.py b/image_layering.py
index 7b1e276..ee8a065 100644
--- a/image_layering.py
+++ b/image_layering.py
@@ -28,7 +28,7 @@ class Layering:
             "optional": {
                 "add_layer2": ("IMAGE", {"default": None}),
                 "add_layer3": ("IMAGE", {"default": None}),
-                "key_color": ("TUPLE", {"default": (255, 255, 255)}),
+                # "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}),
@@ -72,7 +72,7 @@ class Layering:
         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)):
+    def apply_blend(self, base_image, add_layer1, alpha1=1.0, add_layer2=None, alpha2=1.0, add_layer3=None, alpha3=1.0):
         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]