裁剪多余透明边，识别粗线功能

__init__.py

@@ -7,6 +7,7 @@ import torch
 import comfy.utils
 from .videoCut import getCutList, saveToDir
 from .seg import get_masks
+from .thick_lines_from_canny import fill_white_segments
 
 
 def getImageSize(IMAGE) -> tuple[int, int]:
@@ -731,7 +732,87 @@ class deleteDir:
                 return (1, e_info,)
             except Exception as e:
                 e_info = str(e)
-                return (0,e_info,)
+                return (0, e_info,)
+
+
+class FindThickLinesFromCanny:
+    @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 = result.convert("RGB")
+        result_tensor = imgToTensor(result)
+        return (result_tensor,)
+
+
+class TrimTransparentEdges:
+    @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)
+
+        return (cropped_img,)
 
 
 NODE_CLASS_MAPPINGS = {
@@ -752,6 +833,8 @@ NODE_CLASS_MAPPINGS = {
     "CropImageByMask-badger": CropImageByMask,
     "ApplyMaskToImage-badger": ApplyMaskToImage,
     "deleteDir-badger": deleteDir,
+    "FindThickLinesFromCanny-badger": FindThickLinesFromCanny,
+    "TrimTransparentEdges-badger": TrimTransparentEdges
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {

thick_lines_from_canny.py

@@ -0,0 +1,96 @@
+from PIL import Image
+
+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