From 60639b2265d2e3476d1df17994c8574ea2dd455b Mon Sep 17 00:00:00 2001 From: HSDHCdev Date: Thu, 26 Feb 2026 15:48:17 +0500 Subject: [PATCH] feat: add k_means_vibrant_preserved for high contrast outlier colors --- README.md | 34 +-- ai_pixel_art_enhancer.py | 449 +++++++++++++++++++++++++++++++-------- 2 files changed, 384 insertions(+), 99 deletions(-) diff --git a/README.md b/README.md index 61c3a98..a0ee444 100644 --- a/README.md +++ b/README.md @@ -12,11 +12,12 @@ https://civitai.com/models/1631459/pixel-art-style-illustrious-by-skormino - **Multiple Conversion Methods**: Choose from 6 different pixel art conversion algorithms - **AI-Enhanced Processing**: Intelligent noise reduction, edge enhancement, and detail preservation -- **Advanced Color Management**: Color quantization, similarity clustering, and brightness weighting +- **Advanced Color Management**: Color quantization, similarity clustering, brightness weighting, and strict perfect-palette mapping - **Post-Processing Effects**: Dithering, anti-aliasing, and contrast/saturation adjustments - **Comparison Output**: Visual comparison grid showing original, pixel art, and final enhanced result -- **Flexible Grid Sizing**: Support for various pixel art resolutions (8x8 to 256x256) -- **Scalable Output**: Configurable output scaling (1x to 16x) +- **Palette Extraction**: Extracts the precise colors used into a 32x32 unified swatch grid output +- **Flexible Grain Sizing**: Choose exact pixel sizes for your grains while output resolution matches original +- **Seamless Integrability**: No need to worry about strange dimensions, the node maintains input proportions flawlessly ## Installation @@ -58,11 +59,9 @@ https://civitai.com/models/1631459/pixel-art-style-illustrious-by-skormino #### Required Parameters - **image**: Input image tensor -- **grid_width** (8-256): Width of the pixel art grid -- **grid_height** (8-256): Height of the pixel art grid +- **pixel_size** (1-128): Size of each "grain" in pixels. Output resolution will match input exactly. - **conversion_method**: Algorithm used for pixel conversion - **color_similarity_threshold** (5.0-100.0): Threshold for color clustering -- **output_scale** (1-16): Final output scaling factor #### Optional Enhancement Parameters - **enable_ai_enhancement** (Boolean): Enable AI preprocessing and post-processing @@ -74,6 +73,7 @@ https://civitai.com/models/1631459/pixel-art-style-illustrious-by-skormino - **saturation_boost** (0.0-2.0): Color saturation multiplier - **preserve_details** (Boolean): Enable detail preservation during processing - **anti_aliasing** (Boolean): Apply subtle anti-aliasing to final output +- **input_palette** (Image): Provide an image to strictly map the result to its exact colors. Overrides `color_quantization`. ## Conversion Methods @@ -131,16 +131,17 @@ When `enable_ai_enhancement` is true, the node applies a sophisticated processin ## Output -The node provides two outputs: +The node provides three outputs: 1. **Enhanced Image**: The final pixel art result at the specified scale 2. **Comparison Grid**: Side-by-side comparison of original, intermediate, and final images +3. **Palette**: A 32x32 block grid image showing all the exact colors used in the generated pixel art, sorted by hue. ## Tips and Best Practices -### Grid Size Selection -- **8x8 to 16x16**: Extreme pixelation, best for icons or very stylized art -- **32x32 to 64x64**: Classic pixel art resolution, good balance of detail and style -- **128x128+**: High-detail pixel art, maintains more original information +### Grain Size Selection +- **Small grains (1-4)**: High-detail pixel art, maintains more original information, classic modern indie game look +- **Medium grains (8-16)**: Classic retro game aesthetic, good balance of detail and heavily stylized blocks +- **Large grains (32+)**: Extreme abstract pixelation, best for simplistic icons or heavy block-colors ### Method Selection Guide - **Portraits**: Use "average" or "brightness_weighted_light" @@ -149,7 +150,7 @@ The node provides two outputs: - **Artistic images**: Experiment with "brightness_weighted_dark" and dithering ### Performance Optimization -- Larger grid sizes process faster but produce less pixelated results +- Larger grain sizes process faster - Disable AI enhancement for faster processing on simple images - Use color quantization values appropriate to your target (16-64 colors typical) @@ -158,14 +159,14 @@ The node provides two outputs: ### Basic Pixel Art Conversion ``` Load Image → AI Pixel Art Enhancer → Save Image -Settings: 32x32 grid, "most_frequent" method, 4x scale +Settings: pixel_size=8, "most_frequent" method ``` ### Enhanced Artistic Processing ``` Load Image → AI Pixel Art Enhancer → Save Image -Settings: 64x64 grid, "edge_preserving" method, AI enhancement enabled, - 16 colors, 0.3 dithering, 6x scale +Settings: pixel_size=4, "edge_preserving" method, AI enhancement enabled, + 16 colors, 0.3 dithering ``` ## Troubleshooting @@ -194,6 +195,9 @@ This project is licensed under the Apache 2.0. - see the LICENSE file for detail ## Changelog +### v1.1.0 +- Added `k_means_vibrant_preserved` palette extraction method to preserve high-contrast, pure colors (like bright red eyes) during color quantization. + ### v1.0.0 - Initial release with 6 conversion methods diff --git a/ai_pixel_art_enhancer.py b/ai_pixel_art_enhancer.py index 4fd553e..0cf1169 100644 --- a/ai_pixel_art_enhancer.py +++ b/ai_pixel_art_enhancer.py @@ -4,6 +4,7 @@ from PIL import Image, ImageFilter import cv2 from collections import defaultdict import math +import colorsys class AIPixelArtEnhancer: """ @@ -15,97 +16,107 @@ class AIPixelArtEnhancer: def INPUT_TYPES(s): return { "required": { - "image": ("IMAGE",), - "grid_width": ("INT", { - "default": 32, - "min": 8, - "max": 256, + "image": ("IMAGE", {"tooltip": "Input image to convert into pixel art."}), + "pixel_size": ("INT", { + "default": 8, + "min": 1, + "max": 128, "step": 1, - "display": "number" + "display": "number", + "tooltip": "Size of each 'grain' in pixels. Output resolution will match input exactly." }), - "grid_height": ("INT", { - "default": 32, - "min": 8, - "max": 256, - "step": 1, - "display": "number" - }), - "conversion_method": (["most_frequent", "average", "neighbor_aware", "brightness_weighted_light", "brightness_weighted_dark", "edge_preserving"],), + "conversion_method": (["most_frequent", "average", "neighbor_aware", "brightness_weighted_light", "brightness_weighted_dark", "edge_preserving"], {"tooltip": "Algorithm used for conversion. 'most_frequent' (mode), 'average' (mean), 'neighbor_aware' (contextual), 'brightness_weighted_light' (favors brights), 'brightness_weighted_dark' (favors darks), 'edge_preserving' (keeps structural detail)."}), "color_similarity_threshold": ("FLOAT", { "default": 30.0, "min": 5.0, "max": 100.0, "step": 1.0, - "display": "number" - }), - "output_scale": ("INT", { - "default": 4, - "min": 1, - "max": 16, - "step": 1, - "display": "number" + "display": "number", + "tooltip": "Threshold for color clustering. Higher values group more colors together." }), }, "optional": { - "enable_ai_enhancement": ("BOOLEAN", {"default": True}), + "enable_ai_enhancement": ("BOOLEAN", {"default": True, "tooltip": "Enable AI preprocessing and post-processing for better results."}), "noise_reduction": ("FLOAT", { "default": 0.3, "min": 0.0, "max": 1.0, "step": 0.1, - "display": "slider" + "display": "slider", + "tooltip": "Strength of noise reduction filter." }), "edge_enhancement": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1, - "display": "slider" + "display": "slider", + "tooltip": "Intensity of edge enhancement to preserve structural details." }), "color_quantization": ("INT", { "default": 16, - "min": 4, + "min": 2, "max": 256, "step": 1, - "display": "number" + "display": "number", + "tooltip": "Number of colors in the final output. Lower values give a more retro look." }), "dithering_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1, - "display": "slider" + "display": "slider", + "tooltip": "Floyd-Steinberg dithering intensity to smoothly blend colors." }), + "palette_extraction": (["median_cut", "fast_octree", "k_means", "k_means_outlier_preserved", "k_means_vibrant_preserved", "uniform_binning"], {"default": "median_cut", "tooltip": "Method to extract colors. 'median_cut' (fast and balanced), 'fast_octree' (fast for transparency), 'k_means' (accurate 3D spatial), 'k_means_outlier_preserved' (boosts rare high-contrast pixels), 'k_means_vibrant_preserved' (boosts highly saturated/pure colors), 'uniform_binning' (basic grid)."}), + "color_matching_space": (["rgb", "lab", "oklab"], {"default": "rgb", "tooltip": "Color space to use when mapping pixels to the allowed palette. 'rgb' (standard linear math), 'lab' (CIELAB human perception), 'oklab' (newer smooth perception, prevents hue shifts)."}), "contrast_boost": ("FLOAT", { "default": 1.0, "min": 0.5, "max": 2.0, "step": 0.1, - "display": "slider" + "display": "slider", + "tooltip": "Contrast adjustment multiplier." }), "saturation_boost": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, - "display": "slider" + "display": "slider", + "tooltip": "Color saturation multiplier." }), - "preserve_details": ("BOOLEAN", {"default": True}), - "anti_aliasing": ("BOOLEAN", {"default": False}), - "batch_processing": ("BOOLEAN", {"default": True}), + "preserve_details": ("BOOLEAN", {"default": True, "tooltip": "Enable detail preservation during processing."}), + "anti_aliasing": ("BOOLEAN", {"default": False, "tooltip": "Apply subtle anti-aliasing to final output to blend jagged edges."}), + "batch_processing": ("BOOLEAN", {"default": True, "tooltip": "Process entire batches of images (e.g., video frames). Disable for single images."}), + "input_palette": ("IMAGE", {"tooltip": "Provide an image to strictly map the result to its exact colors. Overrides 'color_quantization' maximum colors."}), } } - RETURN_TYPES = ("IMAGE", "IMAGE") - RETURN_NAMES = ("enhanced_image", "comparison_grid") + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("enhanced_image", "comparison_grid", "palette") FUNCTION = "enhance_pixel_art" CATEGORY = "image/enhancement" - def enhance_pixel_art(self, image, grid_width, grid_height, conversion_method, - color_similarity_threshold, output_scale, enable_ai_enhancement=True, + def enhance_pixel_art(self, image, pixel_size, conversion_method, + color_similarity_threshold, enable_ai_enhancement=True, noise_reduction=0.3, edge_enhancement=0.5, color_quantization=16, - dithering_strength=0.0, contrast_boost=1.0, saturation_boost=1.0, - preserve_details=True, anti_aliasing=False, batch_processing=True): + dithering_strength=0.0, palette_extraction="median_cut", color_matching_space="rgb", + contrast_boost=1.0, saturation_boost=1.0, + preserve_details=True, anti_aliasing=False, batch_processing=True, + input_palette=None): + + # Extract custom palette colors if provided + palette_colors = None + if input_palette is not None: + extracted_colors = self.extract_colors_from_tensor(input_palette, color_quantization, method=palette_extraction) + if len(extracted_colors) > 0: + palette_colors = extracted_colors + else: + print("Warning: Provided input_palette was empty or invalid. Ignoring palette input.") + + global_palette = set() # Handle batch processing if len(image.shape) == 4 and batch_processing: @@ -120,16 +131,18 @@ class AIPixelArtEnhancer: frame = image[i] # Get individual frame # Process single frame - enhanced_frame, comparison_frame = self.process_single_frame( - frame, grid_width, grid_height, conversion_method, - color_similarity_threshold, output_scale, enable_ai_enhancement, + enhanced_frame, comparison_frame, frame_colors = self.process_single_frame( + frame, pixel_size, conversion_method, + color_similarity_threshold, enable_ai_enhancement, noise_reduction, edge_enhancement, color_quantization, - dithering_strength, contrast_boost, saturation_boost, - preserve_details, anti_aliasing + dithering_strength, palette_extraction, color_matching_space, + contrast_boost, saturation_boost, preserve_details, anti_aliasing, palette_colors ) enhanced_frames.append(enhanced_frame) comparison_frames.append(comparison_frame) + for c in frame_colors: + global_palette.add(c) # Progress indicator if (i + 1) % 10 == 0 or i == 0: @@ -138,8 +151,9 @@ class AIPixelArtEnhancer: # Stack frames back into batch tensors enhanced_batch = torch.stack(enhanced_frames, dim=0) comparison_batch = torch.stack(comparison_frames, dim=0) + palette_tensor = self.create_palette_image(list(global_palette)) - return (enhanced_batch, comparison_batch) + return (enhanced_batch, comparison_batch, palette_tensor) else: # Single image processing (legacy behavior) @@ -148,31 +162,41 @@ class AIPixelArtEnhancer: else: frame = image - enhanced_frame, comparison_frame = self.process_single_frame( - frame, grid_width, grid_height, conversion_method, - color_similarity_threshold, output_scale, enable_ai_enhancement, + enhanced_frame, comparison_frame, frame_colors = self.process_single_frame( + frame, pixel_size, conversion_method, + color_similarity_threshold, enable_ai_enhancement, noise_reduction, edge_enhancement, color_quantization, - dithering_strength, contrast_boost, saturation_boost, - preserve_details, anti_aliasing + dithering_strength, palette_extraction, color_matching_space, + contrast_boost, saturation_boost, preserve_details, anti_aliasing, palette_colors ) + for c in frame_colors: + global_palette.add(c) + # Add batch dimension for consistency enhanced_batch = enhanced_frame.unsqueeze(0) comparison_batch = comparison_frame.unsqueeze(0) + palette_tensor = self.create_palette_image(list(global_palette)) - return (enhanced_batch, comparison_batch) + return (enhanced_batch, comparison_batch, palette_tensor) - def process_single_frame(self, frame, grid_width, grid_height, conversion_method, - color_similarity_threshold, output_scale, enable_ai_enhancement, + def process_single_frame(self, frame, pixel_size, conversion_method, + color_similarity_threshold, enable_ai_enhancement, noise_reduction, edge_enhancement, color_quantization, - dithering_strength, contrast_boost, saturation_boost, - preserve_details, anti_aliasing): + dithering_strength, palette_extraction, color_matching_space, + contrast_boost, saturation_boost, + preserve_details, anti_aliasing, palette_colors=None): """Process a single frame/image""" # Convert from tensor [H, W, C] to numpy array img_array = (frame.cpu().numpy() * 255).astype(np.uint8) pil_image = Image.fromarray(img_array) + # Calculate grid size based on pixel_size + orig_width, orig_height = pil_image.size + grid_width = max(1, orig_width // pixel_size) + grid_height = max(1, orig_height // pixel_size) + # AI Enhancement preprocessing if enable_ai_enhancement: pil_image = self.ai_preprocess( @@ -193,18 +217,24 @@ class AIPixelArtEnhancer: color_similarity_threshold ) + frame_colors = [] # Post-processing enhancements - if enable_ai_enhancement: - pixel_art = self.post_process_pixel_art( + if enable_ai_enhancement or palette_colors is not None: + pixel_art, frame_colors = self.post_process_pixel_art( pixel_art, color_quantization, dithering_strength, - anti_aliasing + anti_aliasing, + palette_colors, + palette_extraction, + color_matching_space ) + else: + frame_colors = self.extract_colors_from_pil(pixel_art) - # Scale up the result + # Scale up to exactly the original dimensions (maintaining the pixel block size) final_image = pixel_art.resize( - (grid_width * output_scale, grid_height * output_scale), + (orig_width, orig_height), Image.NEAREST ) @@ -215,7 +245,7 @@ class AIPixelArtEnhancer: enhanced_tensor = self.pil_to_tensor_single(final_image) comparison_tensor = self.pil_to_tensor_single(comparison) - return (enhanced_tensor, comparison_tensor) + return (enhanced_tensor, comparison_tensor, frame_colors) def ai_preprocess(self, image, noise_reduction, edge_enhancement, contrast_boost, saturation_boost, preserve_details): """AI-specific preprocessing to improve pixel art conversion quality""" @@ -457,38 +487,69 @@ class AIPixelArtEnhancer: return None - def post_process_pixel_art(self, pixel_art, color_quantization, dithering_strength, anti_aliasing): + def post_process_pixel_art(self, pixel_art, color_quantization, dithering_strength, anti_aliasing, palette_colors=None, palette_extraction="median_cut", color_matching_space="rgb"): """Apply post-processing enhancements""" - # Color quantization - if color_quantization < 256: - pixel_art = self.quantize_colors(pixel_art, color_quantization) + # Color quantization / Palette application + if palette_colors is not None: + pixel_art = self.apply_palette(pixel_art, palette_colors, color_matching_space) + elif color_quantization <= 256: + # When forcing quantization dynamically, we'll still do the matching in the requested space + pixel_art = self.quantize_colors(pixel_art, color_quantization, palette_extraction) + if color_matching_space != "rgb": + # Ensure the quantized image maps strictly using the perceptual space + extracted_colors = self.extract_colors_from_pil(pixel_art) + if len(extracted_colors) > 0: + pixel_art = self.apply_palette(pixel_art, extracted_colors, color_matching_space) + + frame_colors = self.extract_colors_from_pil(pixel_art) # Dithering - if dithering_strength > 0: + if dithering_strength > 0 and palette_colors is None: pixel_art = self.apply_dithering(pixel_art, dithering_strength) # Anti-aliasing (minimal, preserve pixel art aesthetic) if anti_aliasing: pixel_art = self.apply_subtle_antialiasing(pixel_art) - return pixel_art + return pixel_art, frame_colors - def quantize_colors(self, image, num_colors): + def quantize_colors(self, image, num_colors, method="median_cut"): """Reduce the number of colors in the image""" + if method == "uniform_binning": + return self.manual_quantize_rgba(image, num_colors, force_binning=True) + + if method == "k_means": + return self.manual_quantize_rgba(image, num_colors) + + if method == "k_means_outlier_preserved": + return self.manual_quantize_rgba(image, num_colors, enable_outlier_preservation=True) + + if method == "k_means_vibrant_preserved": + return self.manual_quantize_rgba(image, num_colors, enable_vibrant_preservation=True) + + pil_method = Image.MAXCOVERAGE if method == "median_cut" else Image.FASTOCTREE + if image.mode == 'RGBA': - # For RGBA images, use method 2 (Fast Octree) which supports transparency + # Fast Octree works better with transparency mapping, but MAXCOVERAGE works too. try: - return image.quantize(colors=num_colors, method=2).convert('RGBA') + if pil_method == Image.MAXCOVERAGE: + # Pillow doesn't formally support MAXCOVERAGE with RGBA directly without losing alpha sometimes, + # but typically method 2 (FASTOCTREE) is required for full alpha. Let's try what's requested. + return image.quantize(colors=num_colors, method=pil_method).convert('RGBA') + return image.quantize(colors=num_colors, method=pil_method).convert('RGBA') except Exception: - # Fallback: manual quantization for RGBA + # Fallback: K-Means manual quantization return self.manual_quantize_rgba(image, num_colors) else: - # For RGB images, any method works - return image.quantize(colors=num_colors, method=Image.MAXCOVERAGE).convert('RGBA') + # For RGB images, both methods work purely in PIL + return image.quantize(colors=num_colors, method=pil_method).convert('RGBA') - def manual_quantize_rgba(self, image, num_colors): - """Manual color quantization for RGBA images using K-means clustering""" + def manual_quantize_rgba(self, image, num_colors, force_binning=False, enable_outlier_preservation=False, enable_vibrant_preservation=False): + """Manual color quantization using K-means clustering (forces RGBA)""" + if image.mode != 'RGBA': + image = image.convert('RGBA') + img_array = np.array(image) # Separate pixels with alpha > 0 (opaque) and alpha = 0 (transparent) @@ -505,15 +566,63 @@ class AIPixelArtEnhancer: # Use only RGB channels for clustering rgb_pixels = opaque_pixels[:, :3].astype(np.float32) - # Simple K-means clustering (you could use sklearn if available) - try: - from sklearn.cluster import KMeans - kmeans = KMeans(n_clusters=min(num_colors, len(rgb_pixels)), - random_state=42, n_init=10) - cluster_labels = kmeans.fit_predict(rgb_pixels) - cluster_centers = kmeans.cluster_centers_.astype(np.uint8) - except ImportError: - # Fallback: simple binning if sklearn not available + # Simple K-means clustering + if not force_binning: + try: + from sklearn.cluster import KMeans + + # Outlier Preservation Logic + train_pixels = rgb_pixels + if enable_outlier_preservation and len(rgb_pixels) > 0: + # Calculate mean color and distance of every pixel to the mean + mean_color = np.mean(rgb_pixels, axis=0) + distances = np.sum((rgb_pixels - mean_color) ** 2, axis=1) + + # Target top 5% most outlying contrasting pixels + threshold_dist = np.percentile(distances, 95) + outliers = rgb_pixels[distances > threshold_dist] + + # Amplify outlier presence by duplicating them heavily into the training set + # This forces KMeans to allocate a cluster center to them + if len(outliers) > 0: + weight_multiplier = max(1, len(rgb_pixels) // (len(outliers) * 2)) + repeated_outliers = np.repeat(outliers, weight_multiplier, axis=0) + train_pixels = np.vstack([rgb_pixels, repeated_outliers]) + + # Vibrant Preservation Logic + if enable_vibrant_preservation and len(rgb_pixels) > 0: + # Calculate the difference between the max and min RGB channels for each pixel + # Highly vibrant/pure colors will have a large difference (e.g., Red 255 vs Blue 0) + max_channel = np.max(rgb_pixels, axis=1) + min_channel = np.min(rgb_pixels, axis=1) + vibrancy = max_channel - min_channel + + # Target top 5% most vibrant pixels + threshold_vibrancy = np.percentile(vibrancy, 95) + vibrant_pixels = rgb_pixels[vibrancy > threshold_vibrancy] + + # Amplify vibrant pixel presence + if len(vibrant_pixels) > 0: + weight_multiplier = max(1, len(rgb_pixels) // (len(vibrant_pixels) * 2)) + repeated_vibrant = np.repeat(vibrant_pixels, weight_multiplier, axis=0) + train_pixels = np.vstack([rgb_pixels, repeated_vibrant]) + + kmeans = KMeans(n_clusters=min(num_colors, len(rgb_pixels)), + random_state=42, n_init=10) + + kmeans.fit(train_pixels) + cluster_centers = kmeans.cluster_centers_.astype(np.uint8) + + # Assign labels to the original (unweighted) pixels + distances = np.sum((rgb_pixels[:, np.newaxis, :] - cluster_centers[np.newaxis, :, :]) ** 2, axis=2) + cluster_labels = np.argmin(distances, axis=1) + + except ImportError: + force_binning = True + print("sklearn not found. Falling back to uniform binning.") + + if force_binning: + # Fallback: simple uniform binning cluster_centers = [] step = 256 // int(np.cbrt(num_colors)) for r in range(0, 256, step): @@ -623,6 +732,178 @@ class AIPixelArtEnhancer: return tensor + def extract_colors_from_tensor(self, tensor, max_colors, method="median_cut"): + """Extract dominant RGB colors from a tensor using color quantization""" + if tensor is None: + return [] + + try: + # Convert first frame of palette tensor to PIL Image + if len(tensor.shape) == 4: + img_array = (tensor[0].cpu().numpy() * 255).astype(np.uint8) + elif len(tensor.shape) == 3: + img_array = (tensor.cpu().numpy() * 255).astype(np.uint8) + else: + print(f"Warning: Invalid palette tensor shape {tensor.shape}.") + return [] + + pil_img = Image.fromarray(img_array) + + # Quantize to limit the number of colors based on method + quantized = self.quantize_colors(pil_img, max_colors, method) + + # Extract the unique colors from the quantized image + return self.extract_colors_from_pil(quantized) + + except Exception as e: + print(f"Error extracting colors from palette input: {e}") + return [] + + def extract_colors_from_pil(self, image): + """Extract unique RGB colors from a PIL Image""" + img_array = np.array(image) + if len(img_array.shape) == 3: + pixels = img_array.reshape(-1, img_array.shape[2]) + if img_array.shape[2] == 4: + pixels = pixels[pixels[:, 3] > 0] + rgb_pixels = pixels[:, :3] + else: + rgb_pixels = img_array.reshape(-1, 1) + rgb_pixels = np.repeat(rgb_pixels, 3, axis=1) + unique_colors = np.unique(rgb_pixels, axis=0) + return [tuple(c) for c in unique_colors] + + def create_palette_image(self, colors, swatch_size=32): + """Create an image displaying the palette colors in a roughly square grid""" + if not colors: + return torch.zeros((1, swatch_size, swatch_size, 3), dtype=torch.float32) + + # Sort colors somewhat (by hue to keep it visually pleasing) + colors.sort(key=lambda c: colorsys.rgb_to_hsv(c[0]/255.0, c[1]/255.0, c[2]/255.0)[0]) + + num_colors = len(colors) + + # Calculate grid dimensions to be as square as possible + cols = math.ceil(math.sqrt(num_colors)) + rows = math.ceil(num_colors / cols) + + width = cols * swatch_size + height = rows * swatch_size + + palette_img = Image.new('RGB', (width, height), (0, 0, 0)) + from PIL import ImageDraw + draw = ImageDraw.Draw(palette_img) + + for i, color in enumerate(colors): + x = (i % cols) * swatch_size + y = (i // cols) * swatch_size + draw.rectangle([x, y, x + swatch_size, y + swatch_size], fill=tuple(color)) + + return self.pil_to_tensor_single(palette_img).unsqueeze(0) + + def rgb_to_lab(self, rgb_array): + """Convert float32 RGB [0-255] to CIELAB space using OpenCV""" + rgb_img = (rgb_array.reshape(1, -1, 3).astype(np.float32) / 255.0) + lab_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2Lab) + return lab_img.reshape(-1, 3).astype(np.float32) + + def rgb_to_oklab(self, rgb_array): + """Convert float32 RGB [0-255] to OKLAB space""" + rgb = rgb_array.astype(np.float32) / 255.0 + + # sRGB to linear sRGB + mask = rgb > 0.04045 + rgb_lin = np.empty_like(rgb) + rgb_lin[mask] = np.power((rgb[mask] + 0.055) / 1.055, 2.4) + rgb_lin[~mask] = rgb[~mask] / 12.92 + + # linear sRGB to LMS + M1 = np.array([ + [0.4122214708, 0.5363325363, 0.0514459929], + [0.2119034982, 0.6806995451, 0.1073969566], + [0.0883024619, 0.2817188376, 0.6299787005] + ], dtype=np.float32) + + lms = rgb_lin @ M1.T + + # non-linear LMS + lms_ = np.cbrt(np.maximum(lms, 0)) + + # LMS to OKLAB + M2 = np.array([ + [ 0.2104542553, 0.7936177850, -0.0040720468], + [ 1.9779984951, -2.4285922050, 0.4505937099], + [ 0.0259040371, 0.7827717662, -0.8086757660] + ], dtype=np.float32) + + oklab = lms_ @ M2.T + return oklab + + def apply_palette(self, image, palette_colors, color_matching_space="rgb"): + """Map each pixel to the nearest color in the palette.""" + img_array = np.array(image) + height, width = img_array.shape[:2] + pixels = img_array.reshape(-1, img_array.shape[2]) + + if img_array.shape[2] == 4: + opaque_mask = pixels[:, 3] > 0 + rgb_pixels = pixels[opaque_mask, :3] + else: + opaque_mask = np.ones(len(pixels), dtype=bool) + rgb_pixels = pixels[:, :3] + + if len(rgb_pixels) == 0: + return image + + palette_array = np.array(palette_colors, dtype=np.float32) + quantized_rgb = np.zeros_like(rgb_pixels) + + # Convert spaces once for the palette + if color_matching_space == "oklab": + palette_space = self.rgb_to_oklab(palette_array) + elif color_matching_space == "lab": + palette_space = self.rgb_to_lab(palette_array) + else: + palette_space = palette_array + + # Try to use scipy for much faster distance matrix calculation if available + try: + from scipy.spatial.distance import cdist + use_scipy = True + except ImportError: + use_scipy = False + + # Process in chunks to save memory + chunk_size = 10000 + for i in range(0, len(rgb_pixels), chunk_size): + chunk_rgb = rgb_pixels[i:i+chunk_size].astype(np.float32) + + # Convert space for the image chunk + if color_matching_space == "oklab": + chunk_space = self.rgb_to_oklab(chunk_rgb) + elif color_matching_space == "lab": + chunk_space = self.rgb_to_lab(chunk_rgb) + else: + chunk_space = chunk_rgb + + if use_scipy: + # Highly optimized C backend for distance matrices + distances = cdist(chunk_space, palette_space, metric='sqeuclidean') + else: + # Fallback purely to numpy broadcasting + distances = np.sum((chunk_space[:, np.newaxis, :] - palette_space[np.newaxis, :, :]) ** 2, axis=2) + + nearest_idx = np.argmin(distances, axis=1) + quantized_rgb[i:i+chunk_size] = palette_array[nearest_idx] + + if img_array.shape[2] == 4: + pixels[opaque_mask, :3] = quantized_rgb + result_array = pixels.reshape(height, width, 4).astype(np.uint8) + return Image.fromarray(result_array, 'RGBA') + else: + result_array = quantized_rgb.reshape(height, width, 3).astype(np.uint8) + return Image.fromarray(result_array, 'RGB') + # Node registration for ComfyUI NODE_CLASS_MAPPINGS = { "AIPixelArtEnhancer": AIPixelArtEnhancer