217 lines
8.2 KiB
Python
217 lines
8.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
import torch
|
|
from PIL import Image, ImageDraw
|
|
import time
|
|
|
|
class ImageBinarizer:
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "binarize_image"
|
|
CATEGORY = "image/pattern"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"threshold": ("FLOAT", {
|
|
"default": 127,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1,
|
|
"display": "slider"
|
|
})
|
|
}
|
|
}
|
|
|
|
def binarize_image(self, image, threshold):
|
|
if isinstance(image, torch.Tensor):
|
|
img = (image[0] * 255).cpu().numpy().astype(np.uint8)
|
|
else:
|
|
img = (image[0] * 255).astype(np.uint8)
|
|
|
|
if len(img.shape) != 3:
|
|
raise ValueError("输入图像必须是3通道RGB图像")
|
|
height, width, channels = img.shape
|
|
if channels != 3:
|
|
raise ValueError(f"输入图像必须是3通道RGB图像，当前通道数：{channels}")
|
|
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
_, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)
|
|
result_img = Image.fromarray(binary).convert('RGB')
|
|
|
|
result_array = np.array(result_img) / 255.0
|
|
result_tensor = torch.from_numpy(result_array).float()
|
|
result_tensor = result_tensor.unsqueeze(0)
|
|
return (result_tensor,)
|
|
|
|
class CirclePatternProcessor:
|
|
RETURN_TYPES = ("IMAGE", "TUPLE")
|
|
RETURN_NAMES = ("IMAGE", "CIRCLES")
|
|
FUNCTION = "process_circles"
|
|
CATEGORY = "image/pattern"
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"min_radius": ("FLOAT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"max_radius": ("FLOAT", {
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"output_circle_size": ("FLOAT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 50,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"edge_detection_sensitivity": ("FLOAT", {
|
|
"default": 50,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"circle_accumulator_threshold": ("FLOAT", {
|
|
"default": 30,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"min_center_dist": ("FLOAT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"merge_mode": (["keep_first", "use_average"], {
|
|
"default": "keep_first"
|
|
})
|
|
}
|
|
}
|
|
def process_circles(self, image, min_radius, max_radius, output_circle_size, edge_detection_sensitivity, circle_accumulator_threshold, min_center_dist, merge_mode):
|
|
try:
|
|
# 转换图像格式：从Tensor转为NumPy数组
|
|
if isinstance(image, torch.Tensor):
|
|
img = (image[0] * 255).cpu().numpy().astype(np.uint8)
|
|
else:
|
|
img = (image[0] * 255).astype(np.uint8)
|
|
|
|
# 检查图像维度
|
|
if len(img.shape) != 3:
|
|
raise ValueError("输入图像必须是3通道RGB图像")
|
|
height, width, channels = img.shape
|
|
if channels != 3:
|
|
raise ValueError(f"输入图像必须是3通道RGB图像，当前通道数：{channels}")
|
|
|
|
# 参数类型检查和转换
|
|
min_radius_int = max(1, int(min_radius))
|
|
max_radius_int = max(min_radius_int, int(max_radius))
|
|
output_circle_size = max(1, int(output_circle_size))
|
|
edge_detection_sensitivity = float(edge_detection_sensitivity)
|
|
circle_accumulator_threshold = float(circle_accumulator_threshold)
|
|
min_center_dist = max(1, float(min_center_dist))
|
|
|
|
# 转换为灰度图像
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
|
|
circles = cv2.HoughCircles(
|
|
gray,
|
|
cv2.HOUGH_GRADIENT,
|
|
dp=1,
|
|
minDist=min_radius_int,
|
|
param1=edge_detection_sensitivity,
|
|
param2=circle_accumulator_threshold,
|
|
minRadius=min_radius_int,
|
|
maxRadius=max_radius_int
|
|
)
|
|
|
|
# 创建新图像用于绘制标准化的圆
|
|
result_img = Image.new('RGB', (width, height), 'black')
|
|
draw = ImageDraw.Draw(result_img)
|
|
|
|
# 存储检测到的圆形数据
|
|
circle_data = []
|
|
|
|
if circles is not None:
|
|
circles = np.round(circles[0, :]).astype(int)
|
|
# 合并距离较近的圆心
|
|
merged_circles = []
|
|
used_indices = set()
|
|
|
|
for i in range(len(circles)):
|
|
if i in used_indices:
|
|
continue
|
|
|
|
current_circle = circles[i]
|
|
close_circles = [current_circle]
|
|
|
|
# 查找距离当前圆心较近的其他圆心
|
|
for j in range(i + 1, len(circles)):
|
|
if j in used_indices:
|
|
continue
|
|
|
|
other_circle = circles[j]
|
|
dist = np.sqrt((current_circle[0] - other_circle[0])**2 +
|
|
(current_circle[1] - other_circle[1])**2)
|
|
|
|
if dist <= min_center_dist:
|
|
close_circles.append(other_circle)
|
|
used_indices.add(j)
|
|
|
|
# 根据合并模式处理圆心
|
|
if merge_mode == "keep_first":
|
|
merged_circles.append(close_circles[0])
|
|
else: # use_average
|
|
avg_x = int(np.mean([c[0] for c in close_circles]))
|
|
avg_y = int(np.mean([c[1] for c in close_circles]))
|
|
avg_r = int(np.mean([c[2] for c in close_circles]))
|
|
merged_circles.append([avg_x, avg_y, avg_r])
|
|
|
|
used_indices.add(i)
|
|
|
|
# 绘制合并后的圆并存储圆形数据
|
|
for (x, y, r) in merged_circles:
|
|
# 确保圆的坐标在图像范围内
|
|
if 0 <= x < width and 0 <= y < height:
|
|
# 绘制标准大小的白色圆
|
|
draw.ellipse(
|
|
[(x - output_circle_size, y - output_circle_size),
|
|
(x + output_circle_size, y + output_circle_size)],
|
|
fill='white'
|
|
)
|
|
# 存储圆形数据
|
|
circle_data.append((x, y, r))
|
|
|
|
# 转换回 ComfyUI 格式
|
|
result_array = np.array(result_img) / 255.0
|
|
result_tensor = torch.from_numpy(result_array).float()
|
|
result_tensor = result_tensor.unsqueeze(0)
|
|
return (result_tensor, circle_data)
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"处理圆形图案时出错：{str(e)}")
|
|
|
|
# 转换回 ComfyUI 格式
|
|
result_array = np.array(result_img) / 255.0
|
|
result_tensor = torch.from_numpy(result_array).float()
|
|
result_tensor = result_tensor.unsqueeze(0)
|
|
return (result_tensor, circle_data)
|