依靠颜色识别边界

__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,
 
 }
 

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