1291 lines
34 KiB
Python
1291 lines
34 KiB
Python
import math
|
|
import os
|
|
import uuid
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
import comfy.utils
|
|
from .videoCut import getCutList, video_to_frames, cutToDir, frames_to_video
|
|
from .seg import get_masks
|
|
from .line_editor import fill_white_segments, find_largest_white_component
|
|
from .color_editor import get_colors, find_similar_colors, most_common_fuzzy_color, detect_outline
|
|
from .pixel import *
|
|
import gc
|
|
|
|
|
|
def getImageSize(IMAGE) -> tuple[int, int]:
|
|
samples = IMAGE.movedim(-1, 1)
|
|
size = samples.shape[3], samples.shape[2]
|
|
return size
|
|
|
|
|
|
def maskTensorToImgTensor(maskTensor):
|
|
return maskTensor.reshape((-1, 1, maskTensor.shape[-2], maskTensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
|
|
|
|
|
def tensorToImg(imageTensor):
|
|
imaget = imageTensor[0]
|
|
i = 255. * imaget.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
return img
|
|
|
|
|
|
def imgToTensor(img):
|
|
image = np.array(img).astype(np.float32) / 255.0
|
|
imaget = torch.from_numpy(image)[None,]
|
|
return imaget
|
|
|
|
|
|
def img_to_mask(mask):
|
|
mask = mask.convert("RGBA")
|
|
mask = np.array(mask.getchannel('R')).astype(np.float32) / 255.0
|
|
mask = torch.from_numpy(mask)
|
|
mask = mask.unsqueeze(0)
|
|
return mask
|
|
|
|
|
|
def img_to_np(img):
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
img = np.array(img)
|
|
return img
|
|
|
|
|
|
def np_to_img(numpy):
|
|
return Image.fromarray(numpy.astype(np.uint8))
|
|
|
|
|
|
def maskimg_to_mask(mask_img):
|
|
mask = np_to_img(mask_img)
|
|
mask = img_to_mask(mask)
|
|
return mask
|
|
|
|
|
|
def garbage_collect():
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
gc.collect()
|
|
|
|
|
|
class ImageOverlap:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"base_image": ("IMAGE",),
|
|
"additional_image": ("IMAGE",),
|
|
"x": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"y": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
# RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "overlap"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def overlap(self, base_image, additional_image, x, y):
|
|
b_image = tensorToImg(base_image)
|
|
a_image = tensorToImg(additional_image)
|
|
|
|
b_image.paste(a_image, (x, y))
|
|
o_image = imgToTensor(b_image)
|
|
return (o_image,)
|
|
|
|
|
|
class FloatToInt:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"float": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.01,
|
|
"round": 0.01,
|
|
"display": "number"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
# RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "floatToInt"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def floatToInt(self, float):
|
|
return (round(float),)
|
|
|
|
|
|
class IntToString:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"int": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
# RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "intToString"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def intToString(self, int):
|
|
return (str(int),)
|
|
|
|
|
|
class FloatToString:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"float": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.00001,
|
|
"round": False,
|
|
"display": "number"})
|
|
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
# RETURN_NAMES = ("image_output_name",)
|
|
|
|
FUNCTION = "floatToString"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def floatToString(self, float):
|
|
return (str(float),)
|
|
|
|
|
|
class ImageNormalization:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.01,
|
|
"round": 0.01,
|
|
"display": "number"}),
|
|
"height": ("INT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.01,
|
|
"round": 0.01,
|
|
"display": "number"}),
|
|
"target_width": ("INT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.01,
|
|
"round": 0.01,
|
|
"display": "number"}),
|
|
"target_height": ("INT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.01,
|
|
"round": 0.01,
|
|
"display": "number"})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT", "INT", "INT", "INT", "INT",)
|
|
RETURN_NAMES = ("new_width", "new_height", "top", "left", "bottom", "right")
|
|
|
|
FUNCTION = "imageNormalization"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def imageNormalization(self, width, height, target_width, target_height):
|
|
o_ratio = width / height
|
|
ratio = target_width / target_height
|
|
top = 0
|
|
left = 0
|
|
bottom = 0
|
|
right = 0
|
|
nw = 0
|
|
nh = 0
|
|
# 原图比期望尺寸更扁，对齐宽，计算高，补上下
|
|
if (o_ratio >= ratio):
|
|
upratio = target_width / width
|
|
nw = target_width
|
|
nh = round(height * upratio)
|
|
hdiff = target_height - nh
|
|
top = math.floor(hdiff / 2)
|
|
bottom = math.ceil(hdiff / 2)
|
|
else:
|
|
upratio = target_height / height
|
|
nw = round(width * upratio)
|
|
nh = target_height
|
|
wdiff = target_width - nw
|
|
left = math.floor(wdiff / 2)
|
|
right = math.ceil(wdiff / 2)
|
|
|
|
return (nw, nh, top, left, bottom, right,)
|
|
|
|
|
|
class ImageScaleToSide:
|
|
upscale_methods = ["nearest-exact", "bilinear", "area"]
|
|
crop_methods = ["disabled", "center"]
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"side_length": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"side": (["Longest", "Shortest", "Width", "Height"],),
|
|
"upscale_method": (cls.upscale_methods,),
|
|
"crop": (cls.crop_methods,)}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "imageUpscaleToSide"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def imageUpscaleToSide(self, image, upscale_method, side_length: int, side: str, crop):
|
|
samples = image.movedim(-1, 1)
|
|
|
|
size = getImageSize(image)
|
|
|
|
width_B = int(size[0])
|
|
height_B = int(size[1])
|
|
|
|
width = width_B
|
|
height = height_B
|
|
|
|
def determineSide(_side: str) -> tuple[int, int]:
|
|
width, height = 0, 0
|
|
if _side == "Width":
|
|
heigh_ratio = height_B / width_B
|
|
width = side_length
|
|
height = heigh_ratio * width
|
|
elif _side == "Height":
|
|
width_ratio = width_B / height_B
|
|
height = side_length
|
|
width = width_ratio * height
|
|
return width, height
|
|
|
|
if side == "Longest":
|
|
if width > height:
|
|
width, height = determineSide("Width")
|
|
else:
|
|
width, height = determineSide("Height")
|
|
elif side == "Shortest":
|
|
if width < height:
|
|
width, height = determineSide("Width")
|
|
else:
|
|
width, height = determineSide("Height")
|
|
else:
|
|
width, height = determineSide(side)
|
|
|
|
width = math.ceil(width)
|
|
height = math.ceil(height)
|
|
|
|
cls = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
|
|
cls = cls.movedim(1, -1)
|
|
return (cls,)
|
|
|
|
|
|
class StringToFizz:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"text": ("STRING", {"multiline": True})}}
|
|
|
|
RETURN_TYPES = ("STRING", "INT",)
|
|
|
|
FUNCTION = "stringToFizz"
|
|
CATEGORY = "badger"
|
|
|
|
def stringToFizz(self, text):
|
|
textA = text.split("\n")
|
|
lines = 0
|
|
outText = ""
|
|
for line in textA:
|
|
if (len(line) > 0):
|
|
line = "\"" + str(lines) + "\":\"" + line + "\",\n"
|
|
lines = lines + 1
|
|
outText = outText + line
|
|
outText = outText[:-2]
|
|
return (outText, lines,)
|
|
|
|
|
|
class TextListToString:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"texts": ("STRING", {"multiline": True})}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
INPUT_IS_LIST = True
|
|
|
|
FUNCTION = "textListToString"
|
|
CATEGORY = "badger"
|
|
|
|
def textListToString(self, texts):
|
|
fullString = ""
|
|
if len(texts) <= 1:
|
|
return (texts,)
|
|
else:
|
|
for text in texts:
|
|
fullString += text + "\n"
|
|
return (fullString,)
|
|
|
|
|
|
class getImageSide:
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"side_choose": (["short", "long"],)}}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "getImageSide"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def getImageSide(self, image, side_choose):
|
|
|
|
size = getImageSize(image)
|
|
|
|
width = int(size[0])
|
|
height = int(size[1])
|
|
|
|
side = 0
|
|
if width > height:
|
|
if side_choose == "short":
|
|
side = height
|
|
else:
|
|
side = width
|
|
else:
|
|
if side_choose == "short":
|
|
side = width
|
|
else:
|
|
side = height
|
|
|
|
return (side,)
|
|
|
|
|
|
class VideoToFrame:
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_path": ("STRING", {"default": None}),
|
|
"save_name": ("STRING", {"default": "temp"}),
|
|
"min_side_length": ("INT", {
|
|
"default": 512,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"frame_rate": ("INT", {
|
|
"default": 24,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "video_to_frame"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def video_to_frame(self, video_path, save_name, min_side_length, frame_rate):
|
|
videoPath = os.path.abspath(video_path)
|
|
imagePath = video_to_frames(videoPath, min_side_length, frame_rate, save_name)
|
|
|
|
return (imagePath,)
|
|
|
|
|
|
class VideoCutFromDir:
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"frame_dir": ("STRING", {"default": None}),
|
|
"min_frame": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"max_frame": ("INT", {
|
|
"default": 240,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "video_cut_from_dir"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def video_cut_from_dir(self, frame_dir, min_frame, max_frame):
|
|
cutList = getCutList(frame_dir, min_frame, max_frame)
|
|
dirPathString = cutToDir(frame_dir, cutList)
|
|
|
|
return (dirPathString,)
|
|
|
|
|
|
class FrameToVideo:
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"frame_dir": ("STRING", {"default": ""}),
|
|
"save_path": ("STRING", {"default": "result.mp4"}),
|
|
"frame_rate": ("INT", {
|
|
"default": 24,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "frame_to_video"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def frame_to_video(self, frame_dir, save_path, frame_rate):
|
|
save_path = os.path.abspath(save_path)
|
|
frames_to_video(frame_dir, frame_rate, save_path)
|
|
|
|
return (save_path,)
|
|
|
|
|
|
class getParentDir:
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"dir_path": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "getParentdir"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def getParentdir(self, dir_path):
|
|
dir_path = os.path.abspath(dir_path)
|
|
parent_path = os.path.dirname(dir_path)
|
|
return (parent_path,)
|
|
|
|
|
|
class mkdir:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"dir_path": ("STRING", {"default": ""}),
|
|
"new_dir": ("STRING", {"default": "newdir"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "mkdir"
|
|
|
|
CATEGORY = "badger"
|
|
|
|
def mkdir(self, dir_path, new_dir):
|
|
dir_path = os.path.abspath(dir_path)
|
|
new_dir_path = os.path.join(dir_path, new_dir)
|
|
if not os.path.exists(new_dir_path):
|
|
os.mkdir(new_dir_path)
|
|
return (new_dir_path,)
|
|
|
|
|
|
class findCenterOfMask:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("FLOAT", "FLOAT",)
|
|
RETURN_NAMES = ("X", "Y",)
|
|
FUNCTION = "find_center_of_mask"
|
|
|
|
def find_center_of_mask(self, mask):
|
|
if mask.dim() == 3:
|
|
mask = mask.squeeze(0) # Remove the channel dimension if it exists
|
|
assert mask.dim() == 2, "Mask must be 2D"
|
|
|
|
# Create grids for x and y coordinates
|
|
h, w = mask.size()
|
|
x_coords = torch.arange(w).float().to(mask.device)
|
|
y_coords = torch.arange(h).float().to(mask.device)
|
|
|
|
# Compute the center of mass (centroid) of the mask
|
|
total_mass = mask.sum()
|
|
if total_mass > 0:
|
|
x_center = (mask.sum(dim=0) * x_coords).sum() / total_mass
|
|
y_center = (mask.sum(dim=1) * y_coords).sum() / total_mass
|
|
else:
|
|
x_center, y_center = torch.tensor(0), torch.tensor(0)
|
|
|
|
# Convert to int
|
|
X = float(x_center.item())
|
|
Y = float(y_center.item())
|
|
garbage_collect()
|
|
return (X, Y,)
|
|
|
|
|
|
class SegmentToMaskByPoint:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"img": ("IMAGE",),
|
|
"X": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"Y": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 4096.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"dilate": ("INT", {
|
|
"default": 15,
|
|
"min": 0,
|
|
"max": 4096.0,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"sam_ckpt": ("SAM_MODEL",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("MASK", "MASK", "MASK",)
|
|
RETURN_NAMES = ("mask0", "mask1", "mask2",)
|
|
FUNCTION = "seg_to_mask_by_point"
|
|
|
|
def seg_to_mask_by_point(self, img, X, Y, dilate, sam_ckpt):
|
|
img = tensorToImg(img)
|
|
img = img_to_np(img)
|
|
latest_coords = [X, Y]
|
|
masks = get_masks(img, latest_coords, dilate, sam_ckpt)
|
|
mask0 = maskimg_to_mask(masks[0])
|
|
mask1 = maskimg_to_mask(masks[1])
|
|
mask2 = maskimg_to_mask(masks[2])
|
|
garbage_collect()
|
|
return (mask0, mask1, mask2,)
|
|
|
|
|
|
class CropImageByMask:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT",)
|
|
RETURN_NAMES = ("cropped_img", "X", "Y",)
|
|
FUNCTION = "crop_image_by_mask"
|
|
|
|
def crop_image_by_mask(self, image, mask):
|
|
# Ensure the mask is binary
|
|
mask = (mask > 0.5).float()
|
|
|
|
# Find the bounding box of the mask
|
|
if mask.sum() == 0:
|
|
raise ValueError("The mask is empty, cannot determine bounding box for cropping.")
|
|
|
|
# Find indices where the mask is nonzero
|
|
nonzero_indices = torch.nonzero(mask.squeeze(0), as_tuple=True)
|
|
topmost = torch.min(nonzero_indices[0])
|
|
leftmost = torch.min(nonzero_indices[1])
|
|
bottommost = torch.max(nonzero_indices[0])
|
|
rightmost = torch.max(nonzero_indices[1])
|
|
|
|
# Crop the image using the bounding box
|
|
cropped_image = image[:, topmost:bottommost + 1, leftmost:rightmost + 1]
|
|
|
|
# Return the cropped image and the top-left coordinates of the bounding box
|
|
X = int(leftmost)
|
|
Y = int(topmost)
|
|
garbage_collect()
|
|
return (cropped_image, X, Y,)
|
|
|
|
|
|
class ApplyMaskToImage:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("rgba_image",)
|
|
FUNCTION = "apply_mask_to_image"
|
|
|
|
def apply_mask_to_image(self, image, mask):
|
|
image = tensorToImg(image)
|
|
mask = maskTensorToImgTensor(mask)
|
|
mask = tensorToImg(mask)
|
|
mask = mask.convert("L")
|
|
|
|
# 将图片转换为RGBA，以便添加透明度通道
|
|
image = image.convert("RGBA")
|
|
|
|
# 分离图片的通道
|
|
r, g, b, a = image.split()
|
|
|
|
# 将蒙版应用为alpha通道
|
|
new_a = Image.composite(a, Image.new('L', mask.size, 0), mask)
|
|
|
|
# 合并图像通道和新的alpha通道
|
|
result_image = Image.merge('RGBA', (r, g, b, new_a))
|
|
garbage_collect()
|
|
return (imgToTensor(result_image),)
|
|
|
|
|
|
class DeleteDir:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"start": ("STRING", {"default": None}),
|
|
"dir_path": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
OUTPUT_NODE = True
|
|
|
|
RETURN_TYPES = ("INT", "STRING",)
|
|
RETURN_NAMES = ("result", "e_info")
|
|
FUNCTION = "delete_dir"
|
|
|
|
def delete_dir(self, start, dir_path):
|
|
e_info = ""
|
|
status = 0
|
|
abs_dir_path = os.path.abspath(dir_path)
|
|
if not os.path.exists(abs_dir_path):
|
|
e_info = "路径不存在"
|
|
else:
|
|
try:
|
|
# 遍历文件夹中的每个文件或子文件夹
|
|
for root, dirs, files in os.walk(abs_dir_path):
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
os.remove(file_path) # 删除文件
|
|
|
|
for folder in dirs:
|
|
folder_path = os.path.join(root, folder)
|
|
os.rmdir(folder_path) # 删除空文件夹
|
|
|
|
os.rmdir(abs_dir_path) # 最后删除根目录
|
|
status = 1
|
|
e_info = "成功删除"
|
|
except Exception as e:
|
|
status = 0
|
|
e_info = str(e)
|
|
|
|
garbage_collect()
|
|
return (status, e_info,)
|
|
|
|
|
|
class FindThickLinesFromCanny:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"low_threshold": ("FLOAT", {
|
|
"default": 0.01,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.001,
|
|
"display": "number"
|
|
}),
|
|
"high_threshold": ("FLOAT", {
|
|
"default": 0.02,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.001,
|
|
"display": "number"
|
|
}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "find_thick_lines_from_canny"
|
|
|
|
def find_thick_lines_from_canny(self, image, low_threshold, high_threshold):
|
|
img = tensorToImg(image)
|
|
result = fill_white_segments(img, low_threshold, high_threshold)
|
|
result = find_largest_white_component(result)
|
|
result = result.convert("RGB")
|
|
result_tensor = imgToTensor(result)
|
|
garbage_collect()
|
|
return (result_tensor,)
|
|
|
|
|
|
class TrimTransparentEdges:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "trim_transparent_edges"
|
|
|
|
def trim_transparent_edges(self, image):
|
|
img = tensorToImg(image)
|
|
img = img.convert("RGBA")
|
|
|
|
# 获取图片数据
|
|
datas = img.getdata()
|
|
|
|
# 获取非透明像素的边界
|
|
non_transparent_pixels = [
|
|
(i % img.width, i // img.width)
|
|
for i, pix in enumerate(datas)
|
|
if pix[3] != 0
|
|
]
|
|
if not non_transparent_pixels:
|
|
raise ValueError("Image is fully transparent")
|
|
|
|
# 获取非透明像素的最小和最大坐标
|
|
x_min = min(x for x, _ in non_transparent_pixels)
|
|
y_min = min(y for _, y in non_transparent_pixels)
|
|
x_max = max(x for x, _ in non_transparent_pixels)
|
|
y_max = max(y for _, y in non_transparent_pixels)
|
|
|
|
# 裁剪图片
|
|
cropped_img = img.crop((x_min, y_min, x_max + 1, y_max + 1))
|
|
cropped_img = imgToTensor(cropped_img)
|
|
garbage_collect()
|
|
return (cropped_img,)
|
|
|
|
|
|
class ExpandImageWithColor:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"top": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"bottom": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"left": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"right": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"color": ("STRING", {"default": None}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "expand_image_with_color"
|
|
|
|
def expand_image_with_color(self, image, top, bottom, left, right, color=None):
|
|
img = tensorToImg(image)
|
|
img = img.convert("RGBA")
|
|
# Determine the new size of the image
|
|
new_width = img.width + left + right
|
|
new_height = img.height + top + bottom
|
|
|
|
# Create a new image with the new size and the given background color
|
|
if color:
|
|
new_img = Image.new("RGBA", (new_width, new_height), color)
|
|
else:
|
|
# Use transparency if no color was provided
|
|
new_img = Image.new("RGBA", (new_width, new_height), (0, 0, 0, 0))
|
|
|
|
# Paste the original image onto the new image
|
|
new_img.paste(img, (left, top), img)
|
|
|
|
result = imgToTensor(new_img)
|
|
garbage_collect()
|
|
return (result,)
|
|
|
|
|
|
class GetUUID:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"append": ("STRING", {"default": ""}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "get_uuid"
|
|
|
|
def get_uuid(self, append, seed):
|
|
result = uuid.uuid4().hex + append
|
|
return (result,)
|
|
|
|
|
|
class GetDirName:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"dir_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "get_dir_name"
|
|
|
|
def get_dir_name(self, dir_path):
|
|
folder_name = os.path.basename(dir_path)
|
|
return (folder_name,)
|
|
|
|
|
|
class GetColorFromBorder:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"detection_width": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"classification_threshold": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "get_color_from_border"
|
|
|
|
def get_color_from_border(self, image, detection_width, classification_threshold):
|
|
pil_img = tensorToImg(image)
|
|
colors = get_colors(pil_img, detection_width)
|
|
color = most_common_fuzzy_color(colors, classification_threshold)
|
|
garbage_collect()
|
|
return (color,)
|
|
|
|
|
|
class IdentifyColorToMask:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"color": ("STRING", {"default": "#ffffff"}),
|
|
"detection_threshold": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK",)
|
|
FUNCTION = "identify_color_to_mask"
|
|
|
|
def identify_color_to_mask(self, image, color, detection_threshold):
|
|
pil_img = tensorToImg(image)
|
|
mask_img = find_similar_colors(pil_img, color, detection_threshold)
|
|
mask_tensor = imgToTensor(mask_img)
|
|
mask = img_to_mask(mask_img)
|
|
garbage_collect()
|
|
return (mask_tensor, mask,)
|
|
|
|
|
|
class IdentifyBorderColorToMask:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"color": ("STRING", {"default": "#ffffff"}),
|
|
"detection_threshold": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 4096,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK",)
|
|
FUNCTION = "identify_border_color_to_mask"
|
|
|
|
def identify_border_color_to_mask(self, image, color, detection_threshold):
|
|
pil_img = tensorToImg(image)
|
|
mask_img = detect_outline(pil_img, color, detection_threshold)
|
|
mask_tensor = imgToTensor(mask_img)
|
|
mask = img_to_mask(mask_img)
|
|
garbage_collect()
|
|
return (mask_tensor, mask,)
|
|
|
|
|
|
class GarbageCollect:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"start": ("STRING", {"default": "start"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "gc_node"
|
|
OUTPUT_NODE = True
|
|
|
|
def gc_node(self, start, seed):
|
|
garbage_collect()
|
|
return (start,)
|
|
|
|
|
|
class ToPixel:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"original_image": ("IMAGE",),
|
|
"threshold": ("INT", {
|
|
"default": 30,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"pix": ("INT", {
|
|
"default": 64,
|
|
"min": 1,
|
|
"max": 256,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"tile_size": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 128,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"color_card": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "badger"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "image_to_pixel"
|
|
|
|
def image_to_pixel(self, original_image, threshold, pix, tile_size, color_card=None):
|
|
|
|
regular_size = tile_size*pix
|
|
|
|
original_image = tensorToImg(original_image)
|
|
|
|
original_pixels = regular_image(original_image,output_size=(regular_size,regular_size))
|
|
|
|
# 创建新的图像，用于存储像素化的结果
|
|
pixelated_image = Image.new('RGB', (pix, pix))
|
|
|
|
if color_card!=None:
|
|
color_card = tensorToImg(color_card)
|
|
color_palette = load_color_card(color_card)
|
|
# 遍历每个8x8的方块
|
|
for i in range(0, regular_size, tile_size):
|
|
for j in range(0, regular_size, tile_size):
|
|
# 获取当前方块
|
|
block = original_pixels[i:i+tile_size, j:j+tile_size].reshape(-1, 3)
|
|
# 找到主要颜色
|
|
dominant_color = find_dominant_color(block, threshold)
|
|
# 匹配到颜色卡中的颜色
|
|
matched_color = match_color_to_palette(dominant_color, color_palette)
|
|
# 将匹配的颜色赋给对应的像素点
|
|
pixelated_image.putpixel((j // tile_size, i // tile_size), matched_color)
|
|
else:
|
|
# 遍历每个8x8的方块
|
|
for i in range(0, regular_size, tile_size):
|
|
for j in range(0, regular_size, tile_size):
|
|
# 获取当前方块
|
|
block = original_pixels[i:i+tile_size, j:j+tile_size].reshape(-1, 3)
|
|
# 找到主要颜色
|
|
dominant_color = find_dominant_color(block, threshold)
|
|
# 将主要颜色的平均值赋给对应的像素点
|
|
pixelated_image.putpixel((j // tile_size, i // tile_size), tuple(dominant_color.astype(int)))
|
|
|
|
pixelated_image = imgToTensor(pixelated_image)
|
|
return (pixelated_image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageOverlap-badger": ImageOverlap,
|
|
"FloatToInt-badger": FloatToInt,
|
|
"IntToString-badger": IntToString,
|
|
"FloatToString-badger": FloatToString,
|
|
"ImageNormalization-badger": ImageNormalization,
|
|
"ImageScaleToSide-badger": ImageScaleToSide,
|
|
"StringToFizz-badger": StringToFizz,
|
|
"TextListToString-badger": TextListToString,
|
|
"getImageSide-badger": getImageSide,
|
|
"VideoCutFromDir-badger": VideoCutFromDir,
|
|
"FrameToVideo-badger": FrameToVideo,
|
|
"VideoToFrame-badger": VideoToFrame,
|
|
"getParentDir-badger": getParentDir,
|
|
"mkdir-badger": mkdir,
|
|
"findCenterOfMask-badger": findCenterOfMask,
|
|
"SegmentToMaskByPoint-badger": SegmentToMaskByPoint,
|
|
"CropImageByMask-badger": CropImageByMask,
|
|
"ApplyMaskToImage-badger": ApplyMaskToImage,
|
|
"deleteDir-badger": DeleteDir,
|
|
"FindThickLinesFromCanny-badger": FindThickLinesFromCanny,
|
|
"TrimTransparentEdges-badger": TrimTransparentEdges,
|
|
"ExpandImageWithColor-badger": ExpandImageWithColor,
|
|
"GetUUID-badger": GetUUID,
|
|
"GetDirName-badger": GetDirName,
|
|
"GetColorFromBorder-badger": GetColorFromBorder,
|
|
"IdentifyColorToMask-badger":IdentifyColorToMask,
|
|
"IdentifyBorderColorToMask-badger":IdentifyBorderColorToMask,
|
|
"GarbageCollect-badger": GarbageCollect,
|
|
"ToPixel-badger": ToPixel
|
|
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
}
|