orion4d-ComfyUI-Image-Effects/geometric/voronoi_node.py

import numpy as np
import torch
import cv2
from scipy.spatial import Voronoi, voronoi_plot_2d

class VoronoiNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "num_seeds": ("INT", {"default": 100, "min": 10, "max": 500, "step": 10}),
                "color_mode": (["average", "center_point", "random"], {"default": "average"}),
                "cell_outline": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "seed_distribution": (["random", "edge_based", "grid"], {"default": "random"}),
                "outline_color": (["black", "white", "adaptive"], {"default": "black"}),
                "outline_thickness": ("INT", {"default": 2, "min": 1, "max": 5, "step": 1}),
                "edge_threshold": ("FLOAT", {"default": 0.3, "min": 0.1, "max": 1.0, "step": 0.05}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_voronoi"
    CATEGORY = "Image Effects"

    def apply_voronoi(self, image, num_seeds, color_mode, cell_outline,
                     seed_distribution="random", outline_color="black", 
                     outline_thickness=2, edge_threshold=0.3):
        if len(image.shape) == 4:
            img_tensor = image[0]
        else:
            img_tensor = image
        
        img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
        h, w, c = img_np.shape
        
        # Générer les graines
        seeds = self._generate_seeds(img_np, num_seeds, seed_distribution, edge_threshold)
        
        # Créer le diagramme de Voronoï
        result = self._create_voronoi_image(img_np, seeds, color_mode, 
                                          cell_outline, outline_color, outline_thickness)
        
        result_tensor = torch.from_numpy(result.astype(np.float32) / 255.0).unsqueeze(0)
        return (result_tensor,)
    
    def _generate_seeds(self, image, num_seeds, distribution, edge_threshold):
        """Générer les graines pour le diagramme de Voronoï"""
        h, w = image.shape[:2]
        seeds = []
        
        if distribution == "random":
            for _ in range(num_seeds):
                x = np.random.randint(0, w)
                y = np.random.randint(0, h)
                seeds.append([x, y])
                
        elif distribution == "edge_based":
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            edges = cv2.Canny(gray, int(edge_threshold * 100), int(edge_threshold * 200))
            edge_points = np.column_stack(np.where(edges > 0))
            
            if len(edge_points) > 0:
                indices = np.random.choice(len(edge_points), 
                                         min(num_seeds // 2, len(edge_points)), 
                                         replace=False)
                for idx in indices:
                    y, x = edge_points[idx]
                    seeds.append([x, y])
            
            # Compléter avec des points aléatoires
            remaining = num_seeds - len(seeds)
            for _ in range(remaining):
                x = np.random.randint(0, w)
                y = np.random.randint(0, h)
                seeds.append([x, y])
                
        elif distribution == "grid":
            grid_size = int(np.sqrt(num_seeds))
            for i in range(grid_size):
                for j in range(grid_size):
                    if len(seeds) >= num_seeds:
                        break
                    x = int((j + 0.5) * w / grid_size) + np.random.randint(-w//20, w//20)
                    y = int((i + 0.5) * h / grid_size) + np.random.randint(-h//20, h//20)
                    x = np.clip(x, 0, w-1)
                    y = np.clip(y, 0, h-1)
                    seeds.append([x, y])
        
        return np.array(seeds)
    
    def _create_voronoi_image(self, image, seeds, color_mode, outline, outline_color, thickness):
        """Créer l'image avec diagramme de Voronoï"""
        h, w, c = image.shape
        result = np.zeros_like(image)
        
        # Créer une carte de distance pour chaque graine
        for y in range(h):
            for x in range(w):
                # Trouver la graine la plus proche
                distances = np.sum((seeds - np.array([x, y]))**2, axis=1)
                closest_seed_idx = np.argmin(distances)
                closest_seed = seeds[closest_seed_idx]
                
                # Déterminer la couleur de la cellule
                if color_mode == "average":
                    # Couleur moyenne autour de la graine
                    seed_x, seed_y = int(closest_seed[0]), int(closest_seed[1])
                    region_size = 10
                    x1 = max(0, seed_x - region_size)
                    x2 = min(w, seed_x + region_size)
                    y1 = max(0, seed_y - region_size)
                    y2 = min(h, seed_y + region_size)
                    region = image[y1:y2, x1:x2]
                    if region.size > 0:
                        color = np.mean(region.reshape(-1, c), axis=0)
                    else:
                        color = image[seed_y, seed_x]
                        
                elif color_mode == "center_point":
                    # Couleur du point central de la graine
                    seed_x, seed_y = int(closest_seed[0]), int(closest_seed[1])
                    color = image[seed_y, seed_x]
                    
                elif color_mode == "random":
                    # Couleur aléatoire par cellule
                    np.random.seed(closest_seed_idx)
                    color = np.random.randint(0, 256, 3)
                
                result[y, x] = color
        
        # Ajouter les contours si demandé
        if outline:
            outline_img = self._add_voronoi_outlines(result, seeds, outline_color, thickness)
            result = outline_img
        
        return result
    
    def _add_voronoi_outlines(self, image, seeds, outline_color, thickness):
        """Ajouter les contours des cellules de Voronoï"""
        h, w = image.shape[:2]
        
        # Créer une carte des régions
        regions = np.zeros((h, w), dtype=np.int32)
        for y in range(h):
            for x in range(w):
                distances = np.sum((seeds - np.array([x, y]))**2, axis=1)
                regions[y, x] = np.argmin(distances)
        
        # Détecter les frontières
        edges = np.zeros((h, w), dtype=np.uint8)
        for y in range(1, h-1):
            for x in range(1, w-1):
                if (regions[y, x] != regions[y-1, x] or 
                    regions[y, x] != regions[y+1, x] or
                    regions[y, x] != regions[y, x-1] or 
                    regions[y, x] != regions[y, x+1]):
                    edges[y, x] = 255
        
        # Appliquer l'épaisseur
        if thickness > 1:
            kernel = np.ones((thickness, thickness), np.uint8)
            edges = cv2.dilate(edges, kernel, iterations=1)
        
        # Appliquer la couleur de contour
        result = image.copy()
        if outline_color == "black":
            color = [0, 0, 0]
        elif outline_color == "white":
            color = [255, 255, 255]
        elif outline_color == "adaptive":
            # Couleur adaptative basée sur la luminosité locale
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            color = np.where(gray[edges > 0] > 128, [0, 0, 0], [255, 255, 255])
            result[edges > 0] = color
            return result
        
        result[edges > 0] = color
        return result