ComfyUI: Add color blend
This commit is contained in:
parent
d2ee9d24ac
commit
aa45c98ac7
|
@ -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",
|
||||
}
|
||||
|
|
|
@ -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),)
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue
Block a user