diff --git a/color_correct.py b/color_correct.py new file mode 100644 index 0000000..68b7ffc --- /dev/null +++ b/color_correct.py @@ -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, +}