From 9aa934f70ff6adf91efb26aa8e5cb21ec575196a Mon Sep 17 00:00:00 2001 From: melMass Date: Fri, 21 Jul 2023 23:43:46 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=94=A5=20deprecate=20some=20nodes?= =?UTF-8?q?=20and=20fix=20image=20list?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nodes/faceenhance.py | 2 +- nodes/image_processing.py | 172 +++++++++++--------------------------- nodes/mask.py | 49 +++++++---- nodes/number.py | 6 +- nodes/video.py | 4 +- 5 files changed, 85 insertions(+), 148 deletions(-) diff --git a/nodes/faceenhance.py b/nodes/faceenhance.py index 4865413..e2317af 100644 --- a/nodes/faceenhance.py +++ b/nodes/faceenhance.py @@ -44,7 +44,7 @@ class LoadFaceEnhanceModel: [x.name for x in cls.get_models()], {"default": "None"}, ), - "upscale": ("INT", {"default": 2}), + "upscale": ("INT", {"default": 1}), }, "optional": {"bg_upsampler": ("UPSCALE_MODEL", {"default": None})}, } diff --git a/nodes/image_processing.py b/nodes/image_processing.py index 4b3d333..ad302be 100644 --- a/nodes/image_processing.py +++ b/nodes/image_processing.py @@ -90,19 +90,21 @@ class ColorCorrect: @staticmethod def hsv_adjustment(image: torch.Tensor, hue, saturation, value): - image = tensor2pil(image) - hsv_image = image.convert("HSV") + images = tensor2pil(image) + out = [] + for img in images: + hsv_image = img.convert("HSV") - h, s, v = hsv_image.split() + h, s, v = hsv_image.split() - h = h.point(lambda x: (x + hue * 255) % 256) - s = s.point(lambda x: int(x * saturation)) - v = v.point(lambda x: int(x * value)) + h = h.point(lambda x: (x + hue * 255) % 256) + s = s.point(lambda x: int(x * saturation)) + v = v.point(lambda x: int(x * value)) - hsv_image = Image.merge("HSV", (h, s, v)) - rgb_image = hsv_image.convert("RGB") - - return pil2tensor(rgb_image) + hsv_image = Image.merge("HSV", (h, s, v)) + rgb_image = hsv_image.convert("RGB") + out.append(rgb_image) + return pil2tensor(out) @staticmethod def hsv_adjustment_tensor_not_working(image: torch.Tensor, hue, saturation, value): @@ -182,64 +184,6 @@ class ColorCorrect: return (image,) -class HsvToRgb: - """Convert HSV image to RGB""" - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "convert" - CATEGORY = "mtb/image processing" - - def convert(self, image): - image = image.numpy() - - image = image.squeeze() - # image = image.transpose(1,2,3,0) - image = hsv2rgb(image) - image = np.expand_dims(image, axis=0) - - # image = image.transpose(3,0,1,2) - return (torch.from_numpy(image),) - - -class RgbToHsv: - """Convert RGB image to HSV""" - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "convert" - CATEGORY = "mtb/image processing" - - def convert(self, image): - image = image.numpy() - - image = np.squeeze(image) - image = rgb2hsv(image) - image = np.expand_dims(image, axis=0) - - return (torch.from_numpy(image),) - - class ImageCompare: """Compare two images and return a difference image""" @@ -305,37 +249,6 @@ class LoadImageFromUrl: return (pil2tensor(image),) -class Denoise: - """Denoise an image using total variation minimization.""" - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "weight": ( - "FLOAT", - {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}, - ), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "denoise" - CATEGORY = "mtb/image processing" - - def denoise(self, image: torch.Tensor, weight): - image = image.numpy() - image = image.squeeze() - image = denoise_tv_chambolle(image, weight=weight) - - image = np.expand_dims(image, axis=0) - return (torch.from_numpy(image),) - - class Blur: """Blur an image using a Gaussian filter.""" @@ -371,29 +284,29 @@ class Blur: # https://github.com/lllyasviel/AdverseCleaner/blob/main/clean.py -def deglaze_np_img(np_img): - y = np_img.copy() - for _ in range(64): - y = cv2.bilateralFilter(y, 5, 8, 8) - for _ in range(4): - y = guidedFilter(np_img, y, 4, 16) - return y +# def deglaze_np_img(np_img): +# y = np_img.copy() +# for _ in range(64): +# y = cv2.bilateralFilter(y, 5, 8, 8) +# for _ in range(4): +# y = guidedFilter(np_img, y, 4, 16) +# return y -class DeglazeImage: - """Remove adversarial noise from images""" +# class DeglazeImage: +# """Remove adversarial noise from images""" - @classmethod - def INPUT_TYPES(cls): - return {"required": {"image": ("IMAGE",)}} +# @classmethod +# def INPUT_TYPES(cls): +# return {"required": {"image": ("IMAGE",)}} - CATEGORY = "mtb/image processing" +# CATEGORY = "mtb/image processing" - RETURN_TYPES = ("IMAGE",) - FUNCTION = "deglaze_image" +# RETURN_TYPES = ("IMAGE",) +# FUNCTION = "deglaze_image" - def deglaze_image(self, image): - return (np2tensor(deglaze_np_img(tensor2np(image))),) +# def deglaze_image(self, image): +# return (np2tensor(deglaze_np_img(tensor2np(image))),) class MaskToImage: @@ -489,20 +402,32 @@ class ImagePremultiply: def premultiply(self, image, mask, invert): invert = invert == "True" - image = tensor2pil(image) - mask = tensor2pil(mask).convert("L") + images = tensor2pil(image) if invert: - mask = ImageChops.invert(mask) + masks = tensor2pil(mask) # .convert("L") + else: + masks = tensor2pil(1.0 - mask) - image.putalpha(mask) + single = False + if len(mask) == 1: + single = True + + masks = [x.convert("L") for x in masks] + + out = [] + for i, img in enumerate(images): + cur_mask = masks[0] if single else masks[i] + + img.putalpha(cur_mask) + out.append(img) # if invert: # image = Image.composite(image,Image.new("RGBA", image.size, color=(0,0,0,0)), mask) # else: # image = Image.composite(Image.new("RGBA", image.size, color=(0,0,0,0)), image, mask) - return (pil2tensor(image),) + return (pil2tensor(out),) class ImageResizeFactor: @@ -733,12 +658,9 @@ class SaveImageGrid: __nodes__ = [ ColorCorrect, - HsvToRgb, - RgbToHsv, ImageCompare, - Denoise, Blur, - DeglazeImage, + # DeglazeImage, MaskToImage, ColoredImage, ImagePremultiply, diff --git a/nodes/mask.py b/nodes/mask.py index f7f76ab..6745dfe 100644 --- a/nodes/mask.py +++ b/nodes/mask.py @@ -1,6 +1,7 @@ from rembg import remove from ..utils import pil2tensor, tensor2pil from PIL import Image +import comfy.utils class ImageRemoveBackgroundRembg: @@ -65,27 +66,41 @@ class ImageRemoveBackgroundRembg: post_process_mask, bgcolor, ): - image = remove( - data=tensor2pil(image), - alpha_matting=alpha_matting == "True", - alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, - alpha_matting_background_threshold=alpha_matting_background_threshold, - alpha_matting_erode_size=alpha_matting_erode_size, - session=None, - only_mask=False, - post_process_mask=post_process_mask == "True", - bgcolor=None, - ) + pbar = comfy.utils.ProgressBar(image.size(0)) + images = tensor2pil(image) - # extract the alpha to a new image - mask = image.getchannel(3) + out_img = [] + out_mask = [] + out_img_on_bg = [] - # add our bgcolor behind the image - image_on_bg = Image.new("RGBA", image.size, bgcolor) + for img in images: + img_rm = remove( + data=img, + alpha_matting=alpha_matting == "True", + alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, + alpha_matting_background_threshold=alpha_matting_background_threshold, + alpha_matting_erode_size=alpha_matting_erode_size, + session=None, + only_mask=False, + post_process_mask=post_process_mask == "True", + bgcolor=None, + ) - image_on_bg.paste(image, mask=mask) + # extract the alpha to a new image + mask = img_rm.getchannel(3) - return (pil2tensor(image), pil2tensor(mask), pil2tensor(image_on_bg)) + # add our bgcolor behind the image + image_on_bg = Image.new("RGBA", img_rm.size, bgcolor) + + image_on_bg.paste(img_rm, mask=mask) + + out_img.append(img_rm) + out_mask.append(mask) + out_img_on_bg.append(image_on_bg) + + pbar.update(1) + + return (pil2tensor(out_img), pil2tensor(out_mask), pil2tensor(out_img_on_bg)) __nodes__ = [ diff --git a/nodes/number.py b/nodes/number.py index 5113a5a..529c15d 100644 --- a/nodes/number.py +++ b/nodes/number.py @@ -16,7 +16,7 @@ class IntToBool: RETURN_TYPES = ("BOOL",) FUNCTION = "int_to_bool" - CATEGORY = "number" + CATEGORY = "mtb/number" def int_to_bool(self, int): return (bool(int),) @@ -47,7 +47,7 @@ class IntToNumber: RETURN_TYPES = ("NUMBER",) FUNCTION = "int_to_number" - CATEGORY = "number" + CATEGORY = "mtb/number" def int_to_number(self, int): return (int,) @@ -78,7 +78,7 @@ class FloatToNumber: RETURN_TYPES = ("NUMBER",) FUNCTION = "float_to_number" - CATEGORY = "number" + CATEGORY = "mtb/number" def float_to_number(self, float): return (float,) diff --git a/nodes/video.py b/nodes/video.py index 6e9ae67..7e1505c 100644 --- a/nodes/video.py +++ b/nodes/video.py @@ -31,7 +31,7 @@ class LoadImageSequence: } } - CATEGORY = "video" + CATEGORY = "mtb/IO" FUNCTION = "load_image" RETURN_TYPES = ( "IMAGE", @@ -183,7 +183,7 @@ class SaveImageSequence: OUTPUT_NODE = True - CATEGORY = "image" + CATEGORY = "mtb/IO" def save_images( self,