diff --git a/nodes.py b/nodes.py index bfc31b9..60cbbda 100644 --- a/nodes.py +++ b/nodes.py @@ -255,6 +255,60 @@ class BlurMaskFast: return (torch.from_numpy(dup),) +def cv_blur_tensor(images, dx, dy): + if min(dx, dy) > 100: + np_img = torch.nn.functional.interpolate(images.detach().clone().movedim(-1,1), scale_factor=0.1, mode='bilinear').movedim(1,-1).cpu().numpy() + for index, image in enumerate(np_img): + np_img[index] = cv2.GaussianBlur(image, (dx // 20 * 2 + 1, dy // 20 * 2 + 1), 0) + return torch.nn.functional.interpolate(torch.from_numpy(np_img).movedim(-1,1), size=(images.shape[1], images.shape[2]), mode='bilinear').movedim(1,-1) + else: + np_img = images.detach().clone().cpu().numpy() + for index, image in enumerate(np_img): + np_img[index] = cv2.GaussianBlur(image, (dx, dy), 0) + return torch.from_numpy(np_img) + +class ColorMatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "reference": ("IMAGE", ), + "blur": ("INT", {"default": 0, "min": 0, "max": 1023}), + "factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "batch_normalize" + + CATEGORY = "image/filters" + + def batch_normalize(self, images, reference, blur, factor): + t = images.detach().clone() + ref = reference.detach().clone() + if ref.shape[0] < t.shape[0]: + ref = ref[0].unsqueeze(0).repeat(t.shape[0], 1, 1, 1) + + if blur == 0: + mean = torch.mean(t, (1,2), keepdim=True) + mean_ref = torch.mean(ref, (1,2), keepdim=True) + for i in range(t.shape[0]): + for c in range(3): + t[i,:,:,c] /= mean[i,0,0,c] + t[i,:,:,c] *= mean_ref[i,0,0,c] + else: + d = blur * 2 + 1 + blurred = cv_blur_tensor(torch.clamp(t, 0.001, 1), d, d) + blurred_ref = cv_blur_tensor(torch.clamp(ref, 0.001, 1), d, d) + for i in range(t.shape[0]): + for c in range(3): + t[i,:,:,c] /= blurred[i,:,:,c] + t[i,:,:,c] *= blurred_ref[i,:,:,c] + + t = torch.lerp(images, t, factor) + return (t,) + class DilateErodeMask: def __init__(self): pass @@ -449,6 +503,49 @@ class RemapRange: return (torch.from_numpy(i_dup),) +Channel_List = ["red", "green", "blue", "alpha", "white", "black"] +Alpha_List = ["red", "green", "blue", "alpha", "white", "black", "none"] +class ShuffleChannels: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "red": (Channel_List, {"default": "red"}), + "green": (Channel_List, {"default": "green"}), + "blue": (Channel_List, {"default": "blue"}), + "alpha": (Alpha_List, {"default": "none"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "shuffle" + + CATEGORY = "image/filters" + + def shuffle(self, image, red, green, blue, alpha): + ch = 3 if alpha == "none" else 4 + t = torch.zeros((image.shape[0], image.shape[1], image.shape[2], ch), dtype=image.dtype, device=image.device) + image_copy = image.detach().clone() + + ch_key = [red, green, blue, alpha] + for i in range(ch): + if ch_key[i] == "white": + t[:,:,:,i] = 1 + elif ch_key[i] == "red": + t[:,:,:,i] = image_copy[:,:,:,0] + elif ch_key[i] == "green": + t[:,:,:,i] = image_copy[:,:,:,1] + elif ch_key[i] == "blue": + t[:,:,:,i] = image_copy[:,:,:,2] + elif ch_key[i] == "alpha": + if image.shape[3] > 3: + t[:,:,:,i] = image_copy[:,:,:,3] + else: + t[:,:,:,i] = 1 + + return(t,) + class ClampOutliers: def __init__(self): pass @@ -1248,6 +1345,7 @@ NODE_CLASS_MAPPINGS = { "BlurImageFast": BlurImageFast, "BlurMaskFast": BlurMaskFast, "ClampOutliers": ClampOutliers, + "ColorMatch": ColorMatch, "ConvertNormals": ConvertNormals, "DifferenceChecker": DifferenceChecker, "DilateErodeMask": DilateErodeMask, @@ -1262,6 +1360,7 @@ NODE_CLASS_MAPPINGS = { "NormalMapSimple": NormalMapSimple, "OffsetLatentImage": OffsetLatentImage, "RemapRange": RemapRange, + "ShuffleChannels": ShuffleChannels, "Tonemap": Tonemap, "UnJitterImage": UnJitterImage, "UnTonemap": UnTonemap, @@ -1280,6 +1379,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "BlurImageFast": "Blur Image (Fast)", "BlurMaskFast": "Blur Mask (Fast)", "ClampOutliers": "Clamp Outliers", + "ColorMatch": "Color Match", "ConvertNormals": "Convert Normals", "DifferenceChecker": "Difference Checker", "DilateErodeMask": "Dilate/Erode Mask", @@ -1294,6 +1394,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "NormalMapSimple": "Normal Map (Simple)", "OffsetLatentImage": "Offset Latent Image", "RemapRange": "Remap Range", + "ShuffleChannels": "Shuffle Channels", "Tonemap": "Tonemap", "UnJitterImage": "Un-Jitter Image", "UnTonemap": "UnTonemap",