ComfyUI: Add color blend

This commit is contained in:
MMaker 2023-11-15 20:32:42 -05:00
parent d2ee9d24ac
commit aa45c98ac7
Signed by: mmaker
GPG Key ID: CCE79B8FEDA40FB2
3 changed files with 42 additions and 1 deletions

View File

@ -2,8 +2,10 @@ from . import mmaker_color_enhance_comfyui
NODE_CLASS_MAPPINGS = {
"MMakerColorEnhance": mmaker_color_enhance_comfyui.ColorEnhanceComfyNode,
"MMakerColorBlend": mmaker_color_enhance_comfyui.ColorBlendComfyNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MMakerColorEnhance": "Color Enhance",
"MMakerColorBlend": "Color Blend",
}

View File

@ -1,7 +1,7 @@
import torch
import torchvision.transforms.functional as tf
import torchvision.transforms.v2 as v2
from .mmaker_color_enhance_core import color_enhance
from .mmaker_color_enhance_core import color_enhance, color_blend
class ColorEnhanceComfyNode:
def __init__(self):
@ -30,3 +30,32 @@ class ColorEnhanceComfyNode:
images.append(edited_image)
return (torch.stack(images).permute(0, 2, 3, 1),)
class ColorBlendComfyNode:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"image_blend": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply_color_enhance"
CATEGORY = "postprocessing/Effects"
def apply_color_enhance(self, image: torch.Tensor, image_blend: torch.Tensor):
images = []
image_blend = v2.ToDtype(dtype=torch.uint8, scale=True)(image_blend).squeeze().detach().cpu().numpy()
for img in image:
edited_image = v2.ToDtype(dtype=torch.uint8, scale=True)(img).squeeze()
edited_image = color_blend(edited_image.detach().cpu().numpy(), image_blend)
edited_image = tf.to_tensor(edited_image)
images.append(edited_image)
return (torch.stack(images).permute(0, 2, 3, 1),)

View File

@ -10,5 +10,15 @@ def color_enhance(arr, strength: float = 1) -> Image.Image:
lch[:, :, 1] *= 100/(lerp(100, lch[:, :, 1].max(), strength)) # Normalize chroma component
return Image.fromarray(np.array(skimage.color.lab2rgb(lab=skimage.color.lch2lab(lch=lch)) * 255, dtype=np.uint8))
def color_blend(arr1, arr2) -> Image.Image:
hsv1 = skimage.color.rgb2hsv(rgb=np.array(arr1, dtype=np.uint8))
hsv2 = skimage.color.rgb2hsv(rgb=np.array(arr2, dtype=np.uint8))
hsv1[..., 0] = hsv2[..., 0]
hsv1[..., 1] = hsv2[..., 1]
out = skimage.color.hsv2rgb(hsv1)
return Image.fromarray(np.array(out * 255, dtype=np.uint8))
def lerp(a: float, b: float, t: float) -> float:
return (1 - t) * a + t * b