erosDiffusion-ComfyUI-enricos-nodes/ImageColorSampler.py

import torch
import numpy as np
from PIL import Image
import json
import base64
from io import BytesIO
from server import PromptServer
from comfy_execution.graph import ExecutionBlocker

class ImageColorSampler:
    """
    This node allows clicking on an input image to sample colors,
    creating a color palette from the selected sample points.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "sample_points": ("STRING", {"default": "[]", "multiline": True}),
                "palette_size": ("INT", {"default": 128, "min": 32, "max": 512}),
                "sample_size": ("INT", {"default": 5, "min": 1, "max": 15}),
                "wait_for_input": ("BOOLEAN", {"default": True})
            },
            "hidden": {
                "node_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("palette", "sampled_colors")
    FUNCTION = "create_palette"
    CATEGORY = "image/color"
    
    # Track which nodes are waiting for user input
    waiting_nodes = set()
    
    def tensor_to_base64_image(self, tensor):
        """Convert a torch tensor to a base64 encoded image string"""
        # Convert tensor to numpy and then to PIL image
        img_np = tensor.cpu().numpy().squeeze(0)
        img_np = (img_np * 255).astype(np.uint8)
        img_pil = Image.fromarray(img_np)
        
        # Save to a bytes buffer and convert to base64
        buffered = BytesIO()
        img_pil.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return f"data:image/png;base64,{img_str}"
    
    def create_palette(self, image, sample_points, palette_size=128, sample_size=5, wait_for_input=True, node_id=None):
        """
        Creates a color palette from the sampled points on the image.
        
        Args:
            image: Input image tensor
            sample_points: JSON string of sample points coordinates and colors
            palette_size: Size of the palette image (height in pixels)
            sample_size: Size of sample area (radius) for color averaging
            wait_for_input: Whether to block execution waiting for user input
            node_id: Unique ID of this node instance
        
        Returns:
            Tuple of (palette_image, sampled_colors_json)
        """
        # Parse the sample points
        try:
            points = json.loads(sample_points)
        except json.JSONDecodeError:
            points = []
        
        # Check if this is the initial call or a resumption after user input
        is_initial_call = node_id not in self.waiting_nodes
        
        # For initial call, send image data to the UI for interactive editing
        if is_initial_call and wait_for_input:
            # Convert image tensor to a base64 string for sending to UI
            img_base64 = self.tensor_to_base64_image(image)
            
            # Send image and current points to the UI
            ui_data = {
                "image": img_base64,
                "sample_points": points,
                "sample_size": sample_size,
                "node_id": node_id
            }
            
            # Send message to UI to display the image for interaction
            PromptServer.instance.send_sync("image_sampler_init", {"node": node_id, "data": ui_data})
            
            # Add to waiting nodes and block execution
            self.waiting_nodes.add(node_id)
            
            # Return ExecutionBlocker for both outputs
            return (ExecutionBlocker(None), ExecutionBlocker(None))
        
        # Remove from waiting list if resuming
        if node_id in self.waiting_nodes:
            self.waiting_nodes.remove(node_id)
        
        # Convert image tensor to numpy array
        img_np = image.cpu().numpy().squeeze(0)
        
        # Image dimensions
        height, width, _ = img_np.shape
        
        # If no points, return empty palette
        if not points:
            # Create empty palette
            palette_img = np.zeros((palette_size, palette_size, 3), dtype=np.float32)
            palette_tensor = torch.from_numpy(palette_img)[None, ]
            return (palette_tensor, "[]")
            
        # Calculate colors for each sample point
        sampled_colors = []
        for point in points:
            x = int(point["x"] * width)
            y = int(point["y"] * height)
            
            # Ensure coordinates are within bounds
            x = max(sample_size, min(width - sample_size - 1, x))
            y = max(sample_size, min(height - sample_size - 1, y))
            
            # Sample area - take average color in the sample radius
            sample_area = img_np[y-sample_size:y+sample_size+1, x-sample_size:x+sample_size+1]
            avg_color = np.mean(sample_area, axis=(0, 1))
            
            # Convert to 8-bit RGB
            r, g, b = [int(c * 255) for c in avg_color]
            
            # Add to colors list with position info
            sampled_colors.append({
                "position": {"x": point["x"], "y": point["y"]},
                "color": {"r": r, "g": g, "b": b},
                "hex": f"#{r:02X}{g:02X}{b:02X}"
            })
            
        # Create palette image
        num_colors = len(sampled_colors)
        if num_colors == 0:
            # Create empty palette
            palette_img = np.zeros((palette_size, palette_size, 3), dtype=np.float32)
            palette_tensor = torch.from_numpy(palette_img)[None, ]
            return (palette_tensor, "[]")
            
        # Create palette image - a horizontal strip of colors
        stripe_height = palette_size
        stripe_width = palette_size // num_colors if num_colors > 0 else palette_size
        
        palette_img = np.zeros((stripe_height, palette_size, 3), dtype=np.float32)
        
        for i, color_data in enumerate(sampled_colors):
            color = color_data["color"]
            start_x = i * stripe_width
            end_x = (i + 1) * stripe_width if i < num_colors - 1 else palette_size
            
            # Fill the stripe with the color
            palette_img[:, start_x:end_x, 0] = color["r"] / 255.0
            palette_img[:, start_x:end_x, 1] = color["g"] / 255.0
            palette_img[:, start_x:end_x, 2] = color["b"] / 255.0
        
        # Convert to tensor
        palette_tensor = torch.from_numpy(palette_img)[None, ]
        
        # Return palette image and colors data as JSON
        return (palette_tensor, json.dumps(sampled_colors))

# Add routes for handling the continuation of workflow
@PromptServer.instance.routes.post("/image_sampler/continue")
async def image_sampler_continue(request):
    """Handle when user is done selecting color samples and wants to continue"""
    data = await request.json()
    node_id = data.get("node_id")
    sample_points = data.get("sample_points", "[]")
    
    # Update the sample_points widget value to continue with new points
    PromptServer.instance.send_sync("image_sampler_update", {
        "node": node_id,
        "widget_name": "sample_points",
        "value": json.dumps(sample_points)
    })
    
    return {"status": "success"}