fix: 🔥 deprecate some nodes and fix image list

This commit is contained in:
melMass
2023-07-21 23:43:46 +02:00
parent 91bb95da91
commit 9aa934f70f
5 changed files with 85 additions and 148 deletions
+1 -1
View File
@@ -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})},
}
+47 -125
View File
@@ -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,
+32 -17
View File
@@ -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__ = [
+3 -3
View File
@@ -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,)
+2 -2
View File
@@ -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,