From 87c0409a40dbbf355f6a107e1e5656e2505acc83 Mon Sep 17 00:00:00 2001 From: HSDHCdev Date: Mon, 11 Aug 2025 13:23:45 +0500 Subject: [PATCH] Initial Commit --- __init__.py | 39 +++ ai_pixel_art_enhancer.py | 548 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 587 insertions(+) create mode 100644 __init__.py create mode 100644 ai_pixel_art_enhancer.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..26b0087 --- /dev/null +++ b/__init__.py @@ -0,0 +1,39 @@ +""" +AI Pixel Art Enhancer - ComfyUI Custom Node +============================================ + +A comprehensive pixel art enhancement node for ComfyUI that converts regular images +and AI-generated content into high-quality pixel art with advanced processing options. + +Features: +- Multiple conversion algorithms (most frequent, average, brightness-weighted, etc.) +- AI-specific preprocessing (noise reduction, edge enhancement) +- Post-processing effects (dithering, color quantization, anti-aliasing) +- Flexible grid dimensions and output scaling +- Comparison grid generation + +Author: Custom Node +Version: 1.0.0 +""" + +from .ai_pixel_art_enhancer import AIPixelArtEnhancer + +# Export the node class mappings for ComfyUI +NODE_CLASS_MAPPINGS = { + "AIPixelArtEnhancer": AIPixelArtEnhancer +} + +# Display names for the ComfyUI interface +NODE_DISPLAY_NAME_MAPPINGS = { + "AIPixelArtEnhancer": "🎨 AI Pixel Art Enhancer" +} + +# Optional: Add version info +__version__ = "1.0.0" + +# Optional: Add module metadata +__all__ = [ + "AIPixelArtEnhancer", + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS" +] diff --git a/ai_pixel_art_enhancer.py b/ai_pixel_art_enhancer.py new file mode 100644 index 0000000..c772eb4 --- /dev/null +++ b/ai_pixel_art_enhancer.py @@ -0,0 +1,548 @@ +import torch +import numpy as np +from PIL import Image, ImageFilter +import cv2 +from collections import defaultdict +import math + +class AIPixelArtEnhancer: + """ + ComfyUI Node for enhancing AI-generated pixel art with advanced processing options + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "grid_width": ("INT", { + "default": 32, + "min": 8, + "max": 256, + "step": 1, + "display": "number" + }), + "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"],), + "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" + }), + }, + "optional": { + "enable_ai_enhancement": ("BOOLEAN", {"default": True}), + "noise_reduction": ("FLOAT", { + "default": 0.3, + "min": 0.0, + "max": 1.0, + "step": 0.1, + "display": "slider" + }), + "edge_enhancement": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 2.0, + "step": 0.1, + "display": "slider" + }), + "color_quantization": ("INT", { + "default": 16, + "min": 4, + "max": 256, + "step": 1, + "display": "number" + }), + "dithering_strength": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.1, + "display": "slider" + }), + "contrast_boost": ("FLOAT", { + "default": 1.0, + "min": 0.5, + "max": 2.0, + "step": 0.1, + "display": "slider" + }), + "saturation_boost": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.1, + "display": "slider" + }), + "preserve_details": ("BOOLEAN", {"default": True}), + "anti_aliasing": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("enhanced_image", "comparison_grid") + 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, + 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): + + # Convert ComfyUI tensor to PIL Image + if len(image.shape) == 4: + image = image[0] # Remove batch dimension if present + + # Convert from tensor [H, W, C] to numpy array + img_array = (image.cpu().numpy() * 255).astype(np.uint8) + pil_image = Image.fromarray(img_array) + + # AI Enhancement preprocessing + if enable_ai_enhancement: + pil_image = self.ai_preprocess( + pil_image, + noise_reduction, + edge_enhancement, + contrast_boost, + saturation_boost, + preserve_details + ) + + # Convert to pixel art + pixel_art = self.convert_to_pixel_art( + pil_image, + grid_width, + grid_height, + conversion_method, + color_similarity_threshold + ) + + # Post-processing enhancements + if enable_ai_enhancement: + pixel_art = self.post_process_pixel_art( + pixel_art, + color_quantization, + dithering_strength, + anti_aliasing + ) + + # Scale up the result + final_image = pixel_art.resize( + (grid_width * output_scale, grid_height * output_scale), + Image.NEAREST + ) + + # Create comparison grid + comparison = self.create_comparison_grid(pil_image, pixel_art, final_image) + + # Convert back to ComfyUI tensor format + enhanced_tensor = self.pil_to_tensor(final_image) + comparison_tensor = self.pil_to_tensor(comparison) + + return (enhanced_tensor, comparison_tensor) + + def ai_preprocess(self, image, noise_reduction, edge_enhancement, contrast_boost, saturation_boost, preserve_details): + """AI-specific preprocessing to improve pixel art conversion quality""" + + # Convert to numpy for processing + img_array = np.array(image) + + # Noise reduction using bilateral filter + if noise_reduction > 0: + # Convert to BGR for OpenCV + img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR) + d = int(9 * noise_reduction) + sigma_color = 75 * noise_reduction + sigma_space = 75 * noise_reduction + img_bgr = cv2.bilateralFilter(img_bgr, d, sigma_color, sigma_space) + img_array = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + + # Edge enhancement + if edge_enhancement > 0 and preserve_details: + img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(img_gray, 50, 150) + edges = cv2.dilate(edges, np.ones((2,2), np.uint8), iterations=1) + + # Blend edges back into the image + edges_colored = np.stack([edges, edges, edges], axis=2) + img_array = np.clip( + img_array.astype(np.float32) + edges_colored.astype(np.float32) * edge_enhancement * 0.3, + 0, 255 + ).astype(np.uint8) + + image = Image.fromarray(img_array) + + # Contrast and saturation adjustments + if contrast_boost != 1.0: + from PIL import ImageEnhance + enhancer = ImageEnhance.Contrast(image) + image = enhancer.enhance(contrast_boost) + + if saturation_boost != 1.0: + from PIL import ImageEnhance + enhancer = ImageEnhance.Color(image) + image = enhancer.enhance(saturation_boost) + + return image + + def convert_to_pixel_art(self, image, grid_width, grid_height, method, similarity_threshold): + """Convert image to pixel art using various methods""" + + # Calculate cell dimensions + img_width, img_height = image.size + cell_width = img_width / grid_width + cell_height = img_height / grid_height + + # Create output image + pixel_art = Image.new('RGBA', (grid_width, grid_height), (0, 0, 0, 0)) + pixels = pixel_art.load() + + img_array = np.array(image.convert('RGBA')) + + for j in range(grid_height): + for i in range(grid_width): + # Calculate cell bounds + x1 = int(i * cell_width) + y1 = int(j * cell_height) + x2 = int((i + 1) * cell_width) + y2 = int((j + 1) * cell_height) + + # Extract cell pixels + cell_pixels = img_array[y1:y2, x1:x2] + + # Skip transparent cells + if cell_pixels.size == 0: + continue + + # Calculate representative color based on method + color = self.get_representative_color( + cell_pixels, method, similarity_threshold + ) + + if color is not None: + pixels[i, j] = color + + return pixel_art + + def get_representative_color(self, cell_pixels, method, similarity_threshold): + """Calculate representative color for a cell using various methods""" + + # Flatten the cell pixels and filter out transparent ones + flat_pixels = cell_pixels.reshape(-1, 4) + opaque_pixels = flat_pixels[flat_pixels[:, 3] > 0] # Alpha > 0 + + if len(opaque_pixels) == 0: + return None + + rgb_pixels = opaque_pixels[:, :3] # Drop alpha channel + + if method == "average": + color = np.mean(rgb_pixels, axis=0).astype(int) + return tuple(color.tolist() + [255]) + + elif method == "most_frequent": + return self.get_most_frequent_color(rgb_pixels, similarity_threshold) + + elif method == "neighbor_aware": + # Expand sampling area for better context + return self.get_neighbor_aware_color(rgb_pixels) + + elif method in ["brightness_weighted_light", "brightness_weighted_dark"]: + return self.get_brightness_weighted_color(rgb_pixels, method, similarity_threshold) + + elif method == "edge_preserving": + return self.get_edge_preserving_color(cell_pixels) + + else: + # Fallback to average + color = np.mean(rgb_pixels, axis=0).astype(int) + return tuple(color.tolist() + [255]) + + def get_most_frequent_color(self, pixels, similarity_threshold): + """Find most frequent color with clustering""" + if len(pixels) == 0: + return None + + # Simple clustering based on color similarity + clusters = [] + + for pixel in pixels: + r, g, b = pixel[:3] + + # Find matching cluster + matched_cluster = None + for cluster in clusters: + rep_r, rep_g, rep_b = cluster['representative'][:3] + distance = math.sqrt((r-rep_r)**2 + (g-rep_g)**2 + (b-rep_b)**2) + + if distance <= similarity_threshold: + matched_cluster = cluster + break + + if matched_cluster: + matched_cluster['pixels'].append(pixel) + matched_cluster['count'] += 1 + # Update representative as weighted average + total = matched_cluster['count'] + old_rep = matched_cluster['representative'] + matched_cluster['representative'] = [ + int((old_rep[0] * (total-1) + r) / total), + int((old_rep[1] * (total-1) + g) / total), + int((old_rep[2] * (total-1) + b) / total) + ] + else: + clusters.append({ + 'representative': [int(r), int(g), int(b)], + 'pixels': [pixel], + 'count': 1 + }) + + # Return color from largest cluster + if clusters: + largest_cluster = max(clusters, key=lambda x: x['count']) + return tuple(largest_cluster['representative'] + [255]) + + return None + + def get_neighbor_aware_color(self, pixels): + """Get color using neighbor-aware sampling""" + # For now, use weighted average with more weight on central pixels + if len(pixels) == 0: + return None + + # Simple implementation: return average color + color = np.mean(pixels, axis=0).astype(int) + return tuple(color.tolist() + [255]) + + def get_brightness_weighted_color(self, pixels, method, similarity_threshold): + """Get color weighted by brightness preference""" + if len(pixels) == 0: + return None + + weights = [] + for pixel in pixels: + r, g, b = pixel[:3] + brightness = 0.299 * r + 0.587 * g + 0.114 * b + + if method == "brightness_weighted_light": + weight = brightness / 255.0 + else: # brightness_weighted_dark + weight = (255 - brightness) / 255.0 + + weights.append(0.25 + 0.75 * weight) # Minimum weight of 0.25 + + weights = np.array(weights) + weights = weights / np.sum(weights) # Normalize + + # Calculate weighted average + weighted_color = np.average(pixels, axis=0, weights=weights) + return tuple(weighted_color.astype(int).tolist() + [255]) + + def get_edge_preserving_color(self, cell_pixels): + """Get color while preserving edges""" + # Apply edge detection within the cell + if cell_pixels.size == 0: + return None + + gray = np.mean(cell_pixels[:,:,:3], axis=2).astype(np.uint8) + + # Simple edge detection + if gray.shape[0] > 2 and gray.shape[1] > 2: + edges = cv2.Canny(gray, 50, 150) + if np.sum(edges) > 0: + # If edges detected, use median color to preserve structure + flat_pixels = cell_pixels.reshape(-1, 4) + opaque_pixels = flat_pixels[flat_pixels[:, 3] > 0][:, :3] + if len(opaque_pixels) > 0: + color = np.median(opaque_pixels, axis=0).astype(int) + return tuple(color.tolist() + [255]) + + # Fallback to average + flat_pixels = cell_pixels.reshape(-1, 4) + opaque_pixels = flat_pixels[flat_pixels[:, 3] > 0][:, :3] + if len(opaque_pixels) > 0: + color = np.mean(opaque_pixels, axis=0).astype(int) + return tuple(color.tolist() + [255]) + + return None + + def post_process_pixel_art(self, pixel_art, color_quantization, dithering_strength, anti_aliasing): + """Apply post-processing enhancements""" + + # Color quantization + if color_quantization < 256: + pixel_art = self.quantize_colors(pixel_art, color_quantization) + + # Dithering + if dithering_strength > 0: + 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 + + def quantize_colors(self, image, num_colors): + """Reduce the number of colors in the image""" + if image.mode == 'RGBA': + # For RGBA images, use method 2 (Fast Octree) which supports transparency + try: + return image.quantize(colors=num_colors, method=2).convert('RGBA') + except Exception: + # Fallback: manual quantization for RGBA + 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') + + def manual_quantize_rgba(self, image, num_colors): + """Manual color quantization for RGBA images using K-means clustering""" + img_array = np.array(image) + + # Separate pixels with alpha > 0 (opaque) and alpha = 0 (transparent) + height, width, channels = img_array.shape + pixels = img_array.reshape(-1, 4) + + # Find opaque pixels + opaque_mask = pixels[:, 3] > 0 + opaque_pixels = pixels[opaque_mask] + + if len(opaque_pixels) == 0: + return image # All transparent + + # 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 + cluster_centers = [] + step = 256 // int(np.cbrt(num_colors)) + for r in range(0, 256, step): + for g in range(0, 256, step): + for b in range(0, 256, step): + cluster_centers.append([r, g, b]) + if len(cluster_centers) >= num_colors: + break + if len(cluster_centers) >= num_colors: + break + if len(cluster_centers) >= num_colors: + break + + cluster_centers = np.array(cluster_centers[:num_colors]) + + # Assign each pixel to nearest cluster center + cluster_labels = [] + for pixel in rgb_pixels: + distances = np.sum((cluster_centers - pixel) ** 2, axis=1) + cluster_labels.append(np.argmin(distances)) + cluster_labels = np.array(cluster_labels) + + # Create quantized image + quantized_pixels = pixels.copy() + for i, label in enumerate(cluster_labels): + idx = np.where(opaque_mask)[0][i] + quantized_pixels[idx, :3] = cluster_centers[label] + + quantized_array = quantized_pixels.reshape(height, width, 4).astype(np.uint8) + return Image.fromarray(quantized_array, 'RGBA') + + def apply_dithering(self, image, strength): + """Apply Floyd-Steinberg dithering""" + if strength == 0: + return image + + # Simple dithering implementation + img_array = np.array(image).astype(np.float32) + height, width = img_array.shape[:2] + + for y in range(height - 1): + for x in range(1, width - 1): + old_pixel = img_array[y, x, :3].copy() + new_pixel = np.round(old_pixel / 32) * 32 # Quantize + img_array[y, x, :3] = new_pixel + + quant_error = (old_pixel - new_pixel) * strength + + # Distribute error to neighboring pixels + img_array[y, x + 1, :3] += quant_error * 7/16 + img_array[y + 1, x - 1, :3] += quant_error * 3/16 + img_array[y + 1, x, :3] += quant_error * 5/16 + img_array[y + 1, x + 1, :3] += quant_error * 1/16 + + img_array = np.clip(img_array, 0, 255).astype(np.uint8) + return Image.fromarray(img_array) + + def apply_subtle_antialiasing(self, image): + """Apply very subtle anti-aliasing while preserving pixel art look""" + # Use a very light gaussian blur + return image.filter(ImageFilter.GaussianBlur(radius=0.5)) + + def create_comparison_grid(self, original, pixel_art, final): + """Create a comparison grid showing original, pixel art, and final result""" + + # Resize images to same height for comparison + target_height = 512 + + orig_aspect = original.width / original.height + pixel_aspect = pixel_art.width / pixel_art.height + final_aspect = final.width / final.height + + orig_resized = original.resize((int(target_height * orig_aspect), target_height), Image.LANCZOS) + pixel_resized = pixel_art.resize((int(target_height * pixel_aspect), target_height), Image.NEAREST) + final_resized = final.resize((int(target_height * final_aspect), target_height), Image.NEAREST) + + # Create comparison grid + total_width = orig_resized.width + pixel_resized.width + final_resized.width + comparison = Image.new('RGB', (total_width, target_height), (40, 40, 40)) + + comparison.paste(orig_resized, (0, 0)) + comparison.paste(pixel_resized, (orig_resized.width, 0)) + comparison.paste(final_resized, (orig_resized.width + pixel_resized.width, 0)) + + return comparison + + def pil_to_tensor(self, image): + """Convert PIL Image to ComfyUI tensor format""" + if image.mode != 'RGB': + image = image.convert('RGB') + + img_array = np.array(image).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_array) + + # Add batch dimension + tensor = tensor.unsqueeze(0) + + return tensor + +# Node registration for ComfyUI +NODE_CLASS_MAPPINGS = { + "AIPixelArtEnhancer": AIPixelArtEnhancer +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AIPixelArtEnhancer": "AI Pixel Art Enhancer" +} \ No newline at end of file