nodes/HSDHCdev-ComfyUI-AI-Pixel-Art-Enhancer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add k_means_vibrant_preserved for high contrast outlier colors

Browse Source
This commit is contained in:
HSDHCdev
2026-02-26 15:48:17 +05:00
parent 28414b0464
commit 60639b2265
2 changed files with 384 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+19 -15
View File
@@ -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
ai_pixel_art_enhancer.py
+365 -84
View File
@@ -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