""" ComfyUI Custom Node: Image Color Quantizer ========================================== A preprocessing node that cleans up pixelated diagrams and graphics by reducing colors to 2, 4, 8, or 16 dominant colors. This removes pixelation artifacts and color distortions while maintaining the original appearance of the image. """ import torch import numpy as np from PIL import Image from sklearn.cluster import KMeans, MiniBatchKMeans from collections import Counter class ImageColorQuantizer: """ A ComfyUI node that reduces image colors to clean up pixelated graphics. Uses color quantization to find the dominant colors in an image and maps all pixels to their nearest dominant color, effectively removing noise and pixelation artifacts. """ def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "num_colors": (["2", "4", "8", "16"], { "default": "8" }), "quantization_method": (["kmeans", "median_cut", "most_frequent", "kmeans_minibatch"], { "default": "kmeans" }), "dithering": (["none", "floyd_steinberg"], { "default": "none" }), "color_space": (["RGB", "LAB"], { "default": "LAB" }), }, "optional": { "sample_fraction": ("FLOAT", { "default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1, "display": "slider" }), "smoothing_iterations": ("INT", { "default": 0, "min": 0, "max": 5, "step": 1, "display": "slider" }), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("quantized_image", "color_palette") FUNCTION = "quantize_colors" CATEGORY = "image/preprocessing" def quantize_colors(self, image, num_colors, quantization_method, dithering, color_space, sample_fraction=0.5, smoothing_iterations=0): """ Main function to perform color quantization on the input image. Args: image: Input image tensor [B, H, W, C] with values 0-1 num_colors: Number of colors to reduce to (2, 4, 8, or 16) quantization_method: Algorithm to use for finding dominant colors dithering: Whether to apply dithering (none or floyd_steinberg) color_space: Color space for quantization (RGB or LAB) sample_fraction: Fraction of pixels to sample for clustering smoothing_iterations: Number of smoothing passes to apply Returns: Tuple of (quantized_image, color_palette_visualization) """ n_colors = int(num_colors) batch_results = [] palette_results = [] # Process each image in the batch for i in range(image.shape[0]): # Convert tensor to numpy array (0-255 range) img_np = (image[i].cpu().numpy() * 255).astype(np.uint8) # Get the dominant color palette if quantization_method == "kmeans": palette = self._kmeans_palette(img_np, n_colors, color_space, sample_fraction) elif quantization_method == "kmeans_minibatch": palette = self._kmeans_minibatch_palette(img_np, n_colors, color_space, sample_fraction) elif quantization_method == "median_cut": palette = self._median_cut_palette(img_np, n_colors) elif quantization_method == "most_frequent": palette = self._most_frequent_palette(img_np, n_colors) else: palette = self._kmeans_palette(img_np, n_colors, color_space, sample_fraction) # Apply quantization if dithering == "floyd_steinberg": quantized = self._apply_floyd_steinberg(img_np, palette) else: quantized = self._apply_nearest_color(img_np, palette, color_space) # Apply optional smoothing if smoothing_iterations > 0: quantized = self._apply_smoothing(quantized, palette, smoothing_iterations) # Convert back to tensor format (0-1 range) quantized_tensor = torch.from_numpy(quantized.astype(np.float32) / 255.0) batch_results.append(quantized_tensor) # Create palette visualization palette_vis = self._create_palette_visualization(palette, img_np.shape[1], img_np.shape[0]) palette_tensor = torch.from_numpy(palette_vis.astype(np.float32) / 255.0) palette_results.append(palette_tensor) # Stack results back into batch output_image = torch.stack(batch_results, dim=0) output_palette = torch.stack(palette_results, dim=0) return (output_image, output_palette) def _rgb_to_lab(self, rgb): """Convert RGB to LAB color space for better perceptual distance.""" # Normalize RGB values rgb_norm = rgb.astype(np.float32) / 255.0 # Convert to XYZ mask = rgb_norm > 0.04045 rgb_norm[mask] = ((rgb_norm[mask] + 0.055) / 1.055) ** 2.4 rgb_norm[~mask] = rgb_norm[~mask] / 12.92 rgb_norm *= 100 # RGB to XYZ matrix x = rgb_norm[..., 0] * 0.4124564 + rgb_norm[..., 1] * 0.3575761 + rgb_norm[..., 2] * 0.1804375 y = rgb_norm[..., 0] * 0.2126729 + rgb_norm[..., 1] * 0.7151522 + rgb_norm[..., 2] * 0.0721750 z = rgb_norm[..., 0] * 0.0193339 + rgb_norm[..., 1] * 0.1191920 + rgb_norm[..., 2] * 0.9503041 # Normalize for D65 white point x /= 95.047 y /= 100.000 z /= 108.883 # XYZ to LAB epsilon = 0.008856 kappa = 903.3 fx = np.where(x > epsilon, x ** (1/3), (kappa * x + 16) / 116) fy = np.where(y > epsilon, y ** (1/3), (kappa * y + 16) / 116) fz = np.where(z > epsilon, z ** (1/3), (kappa * z + 16) / 116) L = 116 * fy - 16 a = 500 * (fx - fy) b = 200 * (fy - fz) return np.stack([L, a, b], axis=-1) def _lab_to_rgb(self, lab): """Convert LAB back to RGB color space.""" L, a, b = lab[..., 0], lab[..., 1], lab[..., 2] fy = (L + 16) / 116 fx = a / 500 + fy fz = fy - b / 200 epsilon = 0.008856 kappa = 903.3 xr = np.where(fx ** 3 > epsilon, fx ** 3, (116 * fx - 16) / kappa) yr = np.where(L > kappa * epsilon, ((L + 16) / 116) ** 3, L / kappa) zr = np.where(fz ** 3 > epsilon, fz ** 3, (116 * fz - 16) / kappa) # Multiply by reference white x = xr * 95.047 y = yr * 100.000 z = zr * 108.883 # XYZ to RGB r = x * 3.2404542 / 100 + y * -1.5371385 / 100 + z * -0.4985314 / 100 g = x * -0.9692660 / 100 + y * 1.8760108 / 100 + z * 0.0415560 / 100 b_out = x * 0.0556434 / 100 + y * -0.2040259 / 100 + z * 1.0572252 / 100 rgb = np.stack([r, g, b_out], axis=-1) # Apply gamma correction mask = rgb > 0.0031308 rgb[mask] = 1.055 * (rgb[mask] ** (1/2.4)) - 0.055 rgb[~mask] = 12.92 * rgb[~mask] rgb = np.clip(rgb * 255, 0, 255).astype(np.uint8) return rgb def _kmeans_palette(self, img_np, n_colors, color_space, sample_fraction): """Extract dominant colors using K-means clustering.""" h, w, c = img_np.shape pixels = img_np.reshape(-1, 3) # Sample pixels for faster processing n_samples = int(len(pixels) * sample_fraction) n_samples = max(n_samples, n_colors * 10) # Ensure minimum samples indices = np.random.choice(len(pixels), min(n_samples, len(pixels)), replace=False) sampled_pixels = pixels[indices] # Convert to LAB if requested if color_space == "LAB": sampled_pixels = self._rgb_to_lab(sampled_pixels) # Perform K-means clustering kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10) kmeans.fit(sampled_pixels.astype(np.float32)) palette = kmeans.cluster_centers_ # Convert back to RGB if needed if color_space == "LAB": palette = self._lab_to_rgb(palette) else: palette = np.clip(palette, 0, 255).astype(np.uint8) return palette def _kmeans_minibatch_palette(self, img_np, n_colors, color_space, sample_fraction): """Extract dominant colors using Mini-Batch K-means (faster for large images).""" h, w, c = img_np.shape pixels = img_np.reshape(-1, 3) # Sample pixels n_samples = int(len(pixels) * sample_fraction) n_samples = max(n_samples, n_colors * 10) indices = np.random.choice(len(pixels), min(n_samples, len(pixels)), replace=False) sampled_pixels = pixels[indices] # Convert to LAB if requested if color_space == "LAB": sampled_pixels = self._rgb_to_lab(sampled_pixels) # Perform Mini-Batch K-means kmeans = MiniBatchKMeans(n_clusters=n_colors, random_state=42, batch_size=1024) kmeans.fit(sampled_pixels.astype(np.float32)) palette = kmeans.cluster_centers_ if color_space == "LAB": palette = self._lab_to_rgb(palette) else: palette = np.clip(palette, 0, 255).astype(np.uint8) return palette def _median_cut_palette(self, img_np, n_colors): """Extract dominant colors using the Median Cut algorithm.""" # Use PIL's built-in quantization pil_img = Image.fromarray(img_np) quantized = pil_img.quantize(colors=n_colors, method=Image.Quantize.MEDIANCUT) # Extract palette palette_data = quantized.getpalette()[:n_colors * 3] palette = np.array(palette_data).reshape(-1, 3).astype(np.uint8) return palette def _most_frequent_palette(self, img_np, n_colors): """Extract the most frequently occurring colors.""" h, w, c = img_np.shape pixels = img_np.reshape(-1, 3) # Reduce color precision to group similar colors reduced = (pixels // 8) * 8 # Count color frequencies color_tuples = [tuple(c) for c in reduced] counter = Counter(color_tuples) # Get most common colors most_common = counter.most_common(n_colors) palette = np.array([list(c[0]) for c in most_common], dtype=np.uint8) # If we don't have enough unique colors, pad with existing while len(palette) < n_colors: palette = np.vstack([palette, palette[0:1]]) return palette[:n_colors] def _apply_nearest_color(self, img_np, palette, color_space): """Map each pixel to its nearest color in the palette.""" h, w, c = img_np.shape pixels = img_np.reshape(-1, 3) if color_space == "LAB": # Convert both to LAB for comparison pixels_lab = self._rgb_to_lab(pixels) palette_lab = self._rgb_to_lab(palette) # Find nearest color for each pixel distances = np.zeros((len(pixels), len(palette_lab))) for i, color in enumerate(palette_lab): distances[:, i] = np.sqrt(np.sum((pixels_lab - color) ** 2, axis=1)) else: # Use RGB distance distances = np.zeros((len(pixels), len(palette))) for i, color in enumerate(palette): distances[:, i] = np.sqrt(np.sum((pixels.astype(np.float32) - color.astype(np.float32)) ** 2, axis=1)) nearest_indices = np.argmin(distances, axis=1) quantized_pixels = palette[nearest_indices] return quantized_pixels.reshape(h, w, 3) def _apply_floyd_steinberg(self, img_np, palette): """Apply Floyd-Steinberg dithering with the given palette.""" img_float = img_np.astype(np.float32) h, w, c = img_float.shape output = np.zeros_like(img_np) for y in range(h): for x in range(w): old_pixel = img_float[y, x].copy() # Find nearest palette color distances = np.sqrt(np.sum((palette.astype(np.float32) - old_pixel) ** 2, axis=1)) nearest_idx = np.argmin(distances) new_pixel = palette[nearest_idx].astype(np.float32) output[y, x] = palette[nearest_idx] error = old_pixel - new_pixel # Distribute error to neighboring pixels if x + 1 < w: img_float[y, x + 1] += error * 7 / 16 if y + 1 < h: if x > 0: img_float[y + 1, x - 1] += error * 3 / 16 img_float[y + 1, x] += error * 5 / 16 if x + 1 < w: img_float[y + 1, x + 1] += error * 1 / 16 return output def _apply_smoothing(self, img_np, palette, iterations): """Apply smoothing by reassigning pixels based on their neighborhood.""" result = img_np.copy() h, w, c = result.shape for _ in range(iterations): new_result = result.copy() for y in range(1, h - 1): for x in range(1, w - 1): # Get 3x3 neighborhood neighborhood = result[y-1:y+2, x-1:x+2].reshape(-1, 3) # Find most common color in neighborhood color_tuples = [tuple(c) for c in neighborhood] counter = Counter(color_tuples) most_common = counter.most_common(1)[0][0] new_result[y, x] = most_common result = new_result return result def _create_palette_visualization(self, palette, width, height): """Create a visualization of the color palette.""" n_colors = len(palette) swatch_width = width // n_colors vis = np.zeros((min(50, height // 4), width, 3), dtype=np.uint8) for i, color in enumerate(palette): start_x = i * swatch_width end_x = (i + 1) * swatch_width if i < n_colors - 1 else width vis[:, start_x:end_x] = color return vis class ImageColorQuantizerAdvanced: """ Advanced version with additional options for fine-tuning the quantization. """ def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "num_colors": ("INT", { "default": 8, "min": 2, "max": 256, "step": 1 }), "quantization_method": (["kmeans", "median_cut", "octree", "most_frequent"], { "default": "kmeans" }), }, "optional": { "preserve_colors": ("STRING", { "default": "", "multiline": False, "placeholder": "#FFFFFF,#000000 (hex colors to preserve)" }), "background_color": ("STRING", { "default": "", "multiline": False, "placeholder": "#FFFFFF (optional background)" }), "edge_enhancement": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 2.0, "step": 0.1 }), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("quantized_image",) FUNCTION = "quantize_advanced" CATEGORY = "image/preprocessing" def quantize_advanced(self, image, num_colors, quantization_method, preserve_colors="", background_color="", edge_enhancement=0.0): """ Advanced quantization with additional features. """ batch_results = [] for i in range(image.shape[0]): img_np = (image[i].cpu().numpy() * 255).astype(np.uint8) # Parse preserved colors preserved = [] if preserve_colors.strip(): for color_str in preserve_colors.split(","): color_str = color_str.strip() if color_str.startswith("#") and len(color_str) == 7: r = int(color_str[1:3], 16) g = int(color_str[3:5], 16) b = int(color_str[5:7], 16) preserved.append([r, g, b]) # Apply edge enhancement if requested if edge_enhancement > 0: img_np = self._enhance_edges(img_np, edge_enhancement) # Quantize using PIL (simple and reliable) pil_img = Image.fromarray(img_np) # Adjust number of colors if preserving some effective_colors = max(2, num_colors - len(preserved)) if quantization_method == "octree": quantized = pil_img.quantize(colors=effective_colors, method=Image.Quantize.FASTOCTREE) elif quantization_method == "median_cut": quantized = pil_img.quantize(colors=effective_colors, method=Image.Quantize.MEDIANCUT) else: quantized = pil_img.quantize(colors=effective_colors, method=Image.Quantize.MAXCOVERAGE) result = np.array(quantized.convert("RGB")) # Add preserved colors back if preserved: result = self._add_preserved_colors(img_np, result, preserved) result_tensor = torch.from_numpy(result.astype(np.float32) / 255.0) batch_results.append(result_tensor) return (torch.stack(batch_results, dim=0),) def _enhance_edges(self, img_np, strength): """Enhance edges before quantization to preserve details.""" from scipy import ndimage # Convert to grayscale for edge detection gray = np.mean(img_np.astype(np.float32), axis=2) # Sobel edge detection sobel_x = ndimage.sobel(gray, axis=1) sobel_y = ndimage.sobel(gray, axis=0) edges = np.sqrt(sobel_x**2 + sobel_y**2) edges = edges / edges.max() * strength # Enhance original image at edges enhanced = img_np.astype(np.float32) for c in range(3): enhanced[:, :, c] = np.clip(enhanced[:, :, c] * (1 + edges * 0.5), 0, 255) return enhanced.astype(np.uint8) def _add_preserved_colors(self, original, quantized, preserved_colors): """Map pixels close to preserved colors back to those exact colors.""" result = quantized.copy() for preserved in preserved_colors: preserved = np.array(preserved, dtype=np.uint8) # Find pixels in original that are close to this preserved color distances = np.sqrt(np.sum((original.astype(np.float32) - preserved.astype(np.float32)) ** 2, axis=2)) mask = distances < 30 # Threshold for "close" colors result[mask] = preserved return result