circle_detector.py

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)