diff --git a/__init__.py b/__init__.py index b137278..2e5abbb 100644 --- a/__init__.py +++ b/__init__.py @@ -8,6 +8,7 @@ import comfy.utils from .videoCut import getCutList, video_to_frames, cutToDir, frames_to_video from .seg import get_masks from .thick_lines_from_canny import fill_white_segments, find_largest_white_component +from .remove_line import find_most_frequent_color, find_similar_colors def getImageSize(IMAGE) -> tuple[int, int]: @@ -598,6 +599,8 @@ class mkdir: class findCenterOfMask: + def __init__(self) -> None: + pass @classmethod def INPUT_TYPES(s): @@ -639,6 +642,9 @@ class findCenterOfMask: class SegmentToMaskByPoint: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -688,6 +694,8 @@ class SegmentToMaskByPoint: class CropImageByMask: + def __init__(self) -> None: + pass @classmethod def INPUT_TYPES(s): @@ -729,6 +737,9 @@ class CropImageByMask: class ApplyMaskToImage: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -764,6 +775,9 @@ class ApplyMaskToImage: class DeleteDir: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -807,6 +821,9 @@ class DeleteDir: class FindThickLinesFromCanny: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -844,6 +861,9 @@ class FindThickLinesFromCanny: class TrimTransparentEdges: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -887,6 +907,9 @@ class TrimTransparentEdges: class ExpandImageWithColor: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -954,6 +977,9 @@ class ExpandImageWithColor: class GetUUID: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -974,6 +1000,9 @@ class GetUUID: class GetDirName: + def __init__(self) -> None: + pass + @classmethod def INPUT_TYPES(s): return { @@ -992,6 +1021,47 @@ class GetDirName: return (folder_name,) +class IdentifyLinesBasedOnBorderColor: + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "no_bg_image": ("IMAGE",), + "detection_width": ("INT", { + "default": 1, + "min": 1, + "max": 4096, + "step": 1, + "display": "number" + }), + "color_threshold": ("INT", { + "default": 5, + "min": 1, + "max": 1024, + "step": 1, + "display": "number" + }), + }, + } + + CATEGORY = "badger" + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "identify_lines_based_on_borderColor" + + def identify_lines_based_on_borderColor(self, no_bg_image, detection_width, color_threshold): + img = tensorToImg(no_bg_image) + color_string = find_most_frequent_color(img, detection_width) + line_mask_img = find_similar_colors(img, color_string, color_threshold) + msk_img = imgToTensor(line_mask_img) + mask = img_to_mask(line_mask_img) + mask = mask.unsqueeze(0) + return (msk_img, mask,) + + NODE_CLASS_MAPPINGS = { "ImageOverlap-badger": ImageOverlap, "FloatToInt-badger": FloatToInt, @@ -1017,6 +1087,7 @@ NODE_CLASS_MAPPINGS = { "ExpandImageWithColor-badger": ExpandImageWithColor, "GetUUID-badger": GetUUID, "GetDirName-badger": GetDirName, + "IdentifyLinesBasedOnBorderColor-badger": IdentifyLinesBasedOnBorderColor, } diff --git a/remove_line.py b/remove_line.py new file mode 100644 index 0000000..2eb3adf --- /dev/null +++ b/remove_line.py @@ -0,0 +1,65 @@ +from PIL import Image + + +def find_most_frequent_color(PIL_img, n): + img = PIL_img.convert('RGBA') # 确保图片是RGBA模式 + + # 初始化颜色统计字典 + color_count = {} + + # 获取图片尺寸 + width, height = img.size + + for y in range(height): + count = 0 + # 从左到右扫描 + for x in range(width): + r, g, b, a = img.getpixel((x, y)) + if a == 255: + count += 1 + if count <= n: + color = (r, g, b) + color_count[color] = color_count.get(color, 0) + 1 + else: + count = 0 + count = 0 + # 从右到左扫描 + for x in range(width - 1, -1, -1): + r, g, b, a = img.getpixel((x, y)) + if a == 255: + count += 1 + if count <= n: + color = (r, g, b) + color_count[color] = color_count.get(color, 0) + 1 + else: + count = 0 + + # 找到出现次数最多的颜色 + most_frequent_color = max(color_count, key=color_count.get) + + # 返回像素点最多的颜色对应的字符串(格式化为十六进制) + return '#{:02x}{:02x}{:02x}'.format(*most_frequent_color) + + +def is_color_similar(color1, color2, threshold): + return all(abs(c1 - c2) <= threshold for c1, c2 in zip(color1, color2)) + + +def find_similar_colors(image, color_string, threshold): + color_string = color_string[1:] + # 转换颜色字符串为RGB元组 + target_color = tuple(int(color_string[i:i + 2], 16) for i in (0, 2, 4)) + + # 创建一个同样大小的黑色背景图像 + output_image = Image.new('RGB', image.size, (0, 0, 0)) + pixels = image.load() + output_pixels = output_image.load() + + # 遍历每个像素点,检查颜色是否接近目标颜色 + for x in range(image.width): + for y in range(image.height): + if is_color_similar(pixels[x, y], target_color, threshold): + # 将接近的颜色设置为白色 + output_pixels[x, y] = (255, 255, 255) + + return output_image