chore: 将仓库迁移到子目录
This commit is contained in:
DJ-PM
@@ -0,0 +1,216 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user