From 564965104b238f3213d988134e795fcdea77f286 Mon Sep 17 00:00:00 2001
From: Fictiverse <111762798+Fictiverse@users.noreply.github.com>
Date: Sun, 29 Oct 2023 22:28:56 +0100
Subject: [PATCH] Add files via upload

---
 nodes/FV_ColorCorrection.py | 60 +++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 16 deletions(-)

diff --git a/nodes/FV_ColorCorrection.py b/nodes/FV_ColorCorrection.py
index 1133d10..178d737 100644
--- a/nodes/FV_ColorCorrection.py
+++ b/nodes/FV_ColorCorrection.py
@@ -1,37 +1,65 @@
 import cv2
 import numpy as np
+from skimage.exposure import match_histograms
 from PIL import Image
-from skimage import exposure
+from enum import Enum
+import torch
+
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# Define 'BlendType' and 'blendLayers' as needed
+def blendLayers(image1, image2):
+    # Extract the luminance channel from both images
+    image1_luminance = image1.convert("L")
+    image2_luminance = image2.convert("L")
+
+    # Combine the luminance channel from image1 with the color channels of image2
+    r, g, b = image2.split()
+    blended_image = Image.merge("RGB", [image1_luminance, g, b])
+
+    return blended_image
+
 
 class ColorCorrection:
-    """
-    This node provides a simple interface to apply PixelSort blur to the output image.
-    """
     def __init__(self):
         pass
-    
+
     @classmethod
-    def INPUT_TYPES(cls):  # Use cls instead of s
+    def INPUT_TYPES(cls):
         return {
             "required": {
-                "image": ("IMAGE",),
                 "original_image": ("IMAGE",),
+                "correction": ("IMAGE",),  # Add this line
             },
         }
-    
+
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "color_correction"
     CATEGORY = "Fictiverse"
 
-    def color_correction(self, image, original_image):
-        correction_target = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2LAB)
-        corrected_image = cv2.cvtColor(exposure.match_histograms(
-            cv2.cvtColor(np.asarray(original_image), cv2.COLOR_RGB2LAB),
-            correction_target,
-            channel_axis=2
-        ), cv2.COLOR_LAB2RGB).astype("uint8")
-        return (corrected_image,)
 
+        
+        
+    def color_correction(self, original_image, correction):
+        
+        pil_original_image = np.array(tensor2pil(original_image))
+        pil_correction = np.array(tensor2pil(correction))
+        
+        # Perform color correction
+        original_lab = cv2.cvtColor(pil_original_image, cv2.COLOR_RGB2LAB)
+        corrected_lab = cv2.cvtColor(pil_correction, cv2.COLOR_RGB2LAB)
+        corrected_image = cv2.cvtColor(match_histograms(original_lab, corrected_lab, channel_axis=2), cv2.COLOR_LAB2RGB).astype("uint8")
+
+        # You need to implement the blendLayers function and BlendType.LUMINOSITY.
+        # The following line assumes that it exists in your code:
+        image = blendLayers(Image.fromarray(corrected_image), Image.fromarray(pil_original_image))
+        
+        #Return image is in wrong type ?
+        return (image,)
+
+    
 NODE_CLASS_MAPPINGS = {
     "Color correction": ColorCorrection
 }
\ No newline at end of file