命名规格化，GC测试

__init__.py

@@ -7,8 +7,9 @@ import torch
 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 get_colors, find_similar_colors, most_common_fuzzy_color
+from .line_editor import fill_white_segments, find_largest_white_component
+from .color_editor import get_colors, find_similar_colors, most_common_fuzzy_color
+import gc
 
 
 def getImageSize(IMAGE) -> tuple[int, int]:
@@ -17,6 +18,10 @@ def getImageSize(IMAGE) -> tuple[int, int]:
     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()
@@ -757,6 +762,7 @@ class ApplyMaskToImage:
 
     def apply_mask_to_image(self, image, mask):
         image = tensorToImg(image)
+        mask = maskTensorToImgTensor(mask)
         mask = tensorToImg(mask)
         mask = mask.convert("L")
 
@@ -1071,6 +1077,26 @@ class IdentifyLinesBasedOnBorderColor:
         return (msk_img, mask,)
 
 
+class GarbageCollect:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "start": ("STRING", {"default": None}),
+            },
+        }
+
+    CATEGORY = "badger"
+    RETURN_TYPES = ()
+    FUNCTION = "garbage_collect"
+    OUTPUT_NODE = True
+    def garbage_collect(self, start):
+        gc.collect()
+
+
 NODE_CLASS_MAPPINGS = {
     "ImageOverlap-badger": ImageOverlap,
     "FloatToInt-badger": FloatToInt,
@@ -1097,6 +1123,7 @@ NODE_CLASS_MAPPINGS = {
     "GetUUID-badger": GetUUID,
     "GetDirName-badger": GetDirName,
     "IdentifyLinesBasedOnBorderColor-badger": IdentifyLinesBasedOnBorderColor,
+    "GarbageCollect-badger": GarbageCollect,
 
 }
 

color_editor.py

@@ -0,0 +1,110 @@
+from collections import defaultdict
+import numpy as np
+from PIL import Image
+
+
+def rgb_to_hex(rgb_colr):
+    return '#{:02x}{:02x}{:02x}'.format(*rgb_colr)
+
+
+def hex_to_rgb(hex_color):
+    hex_color = hex_color.lstrip('#')
+    return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))
+
+
+def get_colors(PIL_img, n):
+    color_list = []
+    img = PIL_img.convert('RGBA')  # 确保图片是RGBA模式
+
+    # 获取图片尺寸
+    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 != 0:
+                count += 1
+                if count <= n:
+                    color = (r, g, b)
+                    color_list.append(rgb_to_hex(color))
+            else:
+                count = 0
+        count = 0
+        # 从右到左扫描
+        for x in range(width - 1, -1, -1):
+            r, g, b, a = img.getpixel((x, y))
+            if a != 0:
+                count += 1
+                if count <= n:
+                    color = (r, g, b)
+                    color_list.append(rgb_to_hex(color))
+            else:
+                count = 0
+
+
+    return color_list
+
+
+def color_distance(c1, c2):
+    (r1, g1, b1) = c1
+    (r2, g2, b2) = c2
+    return np.sqrt((r1 - r2) ** 2 + (g1 - g2) ** 2 + (b1 - b2) ** 2)
+
+
+def average_color(colors):
+    r = int(np.mean([c[0] for c in colors]))
+    g = int(np.mean([c[1] for c in colors]))
+    b = int(np.mean([c[2] for c in colors]))
+    return f"{r:02x}{g:02x}{b:02x}"
+
+
+def fuzzy_color_grouping(colors, threshold):
+    groups = defaultdict(list)
+
+    for color in colors:
+        rgb = hex_to_rgb(color)
+        placed = False
+
+        for group_color in groups:
+            if color_distance(rgb, hex_to_rgb(group_color)) < threshold:
+                groups[group_color].append(rgb)
+                placed = True
+                break
+
+        if not placed:
+            groups[color].append(rgb)
+
+    return groups
+
+
+def most_common_fuzzy_color(colors, threshold):
+    groups = fuzzy_color_grouping(colors, threshold)
+    largest_group = max(groups, key=lambda k: len(groups[k]))
+    return average_color(groups[largest_group])
+
+
+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):
+    # 转换颜色字符串为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
+
+

line_editor.py

@@ -0,0 +1,129 @@
+from PIL import Image
+from collections import deque
+
+def draw_line(pixels, x0, y0, x1, y1):
+    """Draw a white line from (x0, y0) to (x1, y1) on the provided pixels map."""
+    dx = abs(x1 - x0)
+    dy = abs(y1 - y0)
+    sx = 1 if x0 < x1 else -1
+    sy = 1 if y0 < y1 else -1
+    err = dx - dy
+
+    while True:
+        pixels[x0, y0] = 255
+        if x0 == x1 and y0 == y1:
+            break
+        e2 = 2 * err
+        if e2 > -dy:
+            err -= dy
+            x0 += sx
+        if e2 < dx:
+            err += dx
+            y0 += sy
+
+def fill_white_segments(original_image, low_threshold, high_threshold):
+    # Load the original image and convert it to grayscale
+    original_image = original_image.convert('L')
+    original_pixels = original_image.load()
+    width, height = original_image.size
+
+    low_threshold = int(width*low_threshold)
+    high_threshold = int(width*high_threshold)
+
+    # Create a new black image to draw the lines
+    new_image = Image.new('L', (width, height), 0)
+    new_pixels = new_image.load()
+
+    # Scan horizontally
+    for y in range(height):
+        point_a = None
+        for x in range(width):
+            if original_pixels[x, y] == 255:
+                if point_a is None:
+                    point_a = (x, y)
+                else:
+                    if x - point_a[0] < high_threshold and x - point_a[0] > low_threshold :
+                        draw_line(new_pixels, point_a[0], point_a[1], x, y)
+                        point_a = (x, y)
+                    else:
+                        point_a = (x, y)
+
+
+    # Scan vertically
+    for x in range(width):
+        point_a = None
+        for y in range(height):
+            if original_pixels[x, y] == 255:
+                if point_a is None:
+                    point_a = (x, y)
+                else:
+                    if y - point_a[1] < high_threshold and y - point_a[1] > low_threshold:
+                        draw_line(new_pixels, point_a[0], point_a[1], x, y)
+                        point_a = (x, y)
+                    else:
+                        point_a = (x, y)
+
+    # Scan diagonally (top-left to bottom-right)
+    for diag in range(-height + 1, width):
+        point_a = None
+        for y in range(max(-diag, 0), min(width - diag, height)):
+            x = y + diag
+            if original_pixels[x, y] == 255:
+                if point_a is None:
+                    point_a = (x, y)
+                else:
+                    if max(abs(x - point_a[0]), abs(y - point_a[1])) < high_threshold and max(abs(x - point_a[0]), abs(y - point_a[1])) > low_threshold:
+                        draw_line(new_pixels, point_a[0], point_a[1], x, y)
+                        point_a = (x, y)
+                    else:
+                        point_a = (x, y)
+
+    # Scan diagonally (top-right to bottom-left)
+    for diag in range(0, width + height):
+        point_a = None
+        for y in range(max(diag - width + 1, 0), min(diag + 1, height)):
+            x = diag - y
+            if original_pixels[x, y] == 255:
+                if point_a is None:
+                    point_a = (x, y)
+                else:
+                    if max(abs(x - point_a[0]), abs(y - point_a[1])) < high_threshold and max(abs(x - point_a[0]), abs(y - point_a[1])) > low_threshold:
+                        draw_line(new_pixels, point_a[0], point_a[1], x, y)
+                        point_a = (x, y)
+                    else:
+                        point_a = (x, y)
+
+    # Save the new image with only the drawn lines
+    return new_image
+
+def find_largest_white_component(image):
+    width, height = image.size
+    visited = set()
+    largest_component = []
+    largest_size = 0
+
+    def bfs(x, y):
+        queue = deque([(x, y)])
+        local_visited = set()
+        while queue:
+            x, y = queue.popleft()
+            if (x, y) not in visited and 0 <= x < width and 0 <= y < height and image.getpixel((x, y)) == 255:
+                visited.add((x, y))
+                local_visited.add((x, y))
+                queue.extend([(x+1, y), (x-1, y), (x, y+1), (x, y-1)])
+        return local_visited
+
+    for y in range(height):
+        for x in range(width):
+            if image.getpixel((x, y)) == 255 and (x, y) not in visited:
+                component = bfs(x, y)
+                if len(component) > largest_size:
+                    largest_size = len(component)
+                    largest_component = component
+
+    # 创建一个新的图像来绘制最大的白色像素点整体
+    new_image = Image.new('1', image.size)
+    for x, y in largest_component:
+        new_image.putpixel((x, y), 255)
+
+    return new_image