adds color correction node

2023-03-30 15:28:12 -04:00
parent 4ac573b830
commit 82d6271f4b
1 changed files with 112 additions and 0 deletions
@@ -0,0 +1,112 @@
+import numpy as np
+import cv2
+import torch
+from PIL import Image, ImageEnhance
+
+class ColorCorrect:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "temperature": ("FLOAT", {
+                    "default": 0,
+                    "min": -100,
+                    "max": 100,
+                    "step": 5
+                }),
+                "hue": ("FLOAT", {
+                    "default": 0,
+                    "min": -90,
+                    "max": 90,
+                    "step": 5
+                }),
+                "brightness": ("FLOAT", {
+                    "default": 0,
+                    "min": -100,
+                    "max": 100,
+                    "step": 5
+                }),
+                "contrast": ("FLOAT", {
+                    "default": 0,
+                    "min": -100,
+                    "max": 100,
+                    "step": 5
+                }),
+                "saturation": ("FLOAT", {
+                    "default": 0,
+                    "min": -100,
+                    "max": 100,
+                    "step": 5
+                }),
+                "gamma": ("FLOAT", {
+                    "default": 1,
+                    "min": 0.2,
+                    "max": 2.2,
+                    "step": 0.1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "color_correct"
+
+    CATEGORY = "postprocessing"
+
+    def color_correct(self, image, temperature, hue, brightness, contrast, saturation, gamma):
+        tensor_img = image.numpy()[0]
+
+        brightness /= 100
+        contrast /= 100
+        saturation /= 100
+        temperature /= 100
+
+        brightness = 1 + brightness
+        contrast = 1 + contrast
+        saturation = 1 + saturation
+
+
+        modified_image = Image.fromarray((tensor_img * 255).astype(np.uint8))
+
+        # brightness
+        modified_image = ImageEnhance.Brightness(modified_image).enhance(brightness)
+
+        # contrast
+        modified_image = ImageEnhance.Contrast(modified_image).enhance(contrast)
+        modified_image = np.array(modified_image).astype(np.float32)
+
+        # temperature
+        if temperature > 0:
+            modified_image[:, :, 0] *= 1 + temperature
+            modified_image[:, :, 1] *= 1 + temperature * 0.4
+        elif temperature < 0:
+            modified_image[:, :, 2] *= 1 - temperature
+        modified_image = np.clip(modified_image, 0, 255)/255
+
+        # gamma
+        modified_image = np.clip(np.power(modified_image, gamma), 0, 1)
+
+        # saturation
+        hls_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HLS)
+        hls_img[:, :, 2] = np.clip(saturation*hls_img[:, :, 2], 0, 1)
+        modified_image = cv2.cvtColor(hls_img, cv2.COLOR_HLS2RGB) * 255
+
+        # hue
+        hsv_img = cv2.cvtColor(modified_image, cv2.COLOR_RGB2HSV)
+        hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue) % 360
+
+        modified_image = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB)
+
+        modified_image = modified_image.astype(np.uint8)
+
+        modified_image = modified_image / 255
+        modified_image = torch.from_numpy(modified_image).unsqueeze(0)
+
+        return (modified_image, )
+
+NODE_CLASS_MAPPINGS = {
+    "ColorCorrect": ColorCorrect,
+}