commit 45fa7a93ecce49ca2889d340e6c58465bfd8cceb Author: Marco Aurélio G. Da Silva Date: Fri Dec 20 01:31:34 2024 -0300 Fisrt Commit diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..fd5df23 --- /dev/null +++ b/__init__.py @@ -0,0 +1,24 @@ +# custom_nodes/image_processing/__init__.py +from .custom_crop import CustomCropNode +from .smart_resize import SmartResizeNode +from .nearest_upscale import NearestUpscaleNode +from .load_images import LoadImagesOriginalSize +from .pixel_normalizer import PixelArtNormalizerNode + +NODE_CLASS_MAPPINGS = { + "CustomCrop": CustomCropNode, + "SmartResize": SmartResizeNode, + "NearestUpscale": NearestUpscaleNode, + "LoadImagesOriginal": LoadImagesOriginalSize, + "PixelArtNormalizer": PixelArtNormalizerNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CustomCrop": "Custom Crop", + "SmartResize": "Smart Resize with Border Fill", + "NearestUpscale": "Nearest Neighbor Upscale", + "LoadImagesOriginal": "Load Images (Original Size)", + "PixelArtNormalizer": "Pixel Art Normalizer" +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/custom_crop.py b/custom_crop.py new file mode 100644 index 0000000..ecdb9e6 --- /dev/null +++ b/custom_crop.py @@ -0,0 +1,74 @@ +# custom_nodes/image_processing/custom_crop.py +import torch +import numpy as np +from PIL import Image + +class CustomCropNode: + def __init__(self): + self.crop_modes = ["center", "left", "right", "top", "bottom"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "crop_width": ("INT", { + "default": 1024, + "min": 64, + "max": 8192, + "step": 64 + }), + "crop_height": ("INT", { + "default": 1024, + "min": 64, + "max": 8192, + "step": 64 + }), + "crop_mode": (["center", "left", "right", "top", "bottom"],), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "crop_image" + CATEGORY = "image/processing" + + def crop_image(self, image, crop_width, crop_height, crop_mode): + # Converter o tensor para PIL Image para facilitar o cropping + if isinstance(image, torch.Tensor): + image_np = image[0].cpu().numpy() + image_pil = Image.fromarray((image_np * 255).astype(np.uint8)) + else: + image_pil = image + + # Pegar dimensões originais + orig_width, orig_height = image_pil.size + + # Calcular coordenadas de crop baseado no modo + if crop_mode == "center": + left = (orig_width - crop_width) // 2 + top = (orig_height - crop_height) // 2 + elif crop_mode == "left": + left = 0 + top = (orig_height - crop_height) // 2 + elif crop_mode == "right": + left = orig_width - crop_width + top = (orig_height - crop_height) // 2 + elif crop_mode == "top": + left = (orig_width - crop_width) // 2 + top = 0 + else: # bottom + left = (orig_width - crop_width) // 2 + top = orig_height - crop_height + + # Ajustar coordenadas se necessário para evitar crops fora da imagem + left = max(0, min(left, orig_width - crop_width)) + top = max(0, min(top, orig_height - crop_height)) + + # Realizar o crop + cropped_image = image_pil.crop((left, top, left + crop_width, top + crop_height)) + + # Converter de volta para tensor + cropped_np = np.array(cropped_image).astype(np.float32) / 255.0 + cropped_tensor = torch.from_numpy(cropped_np).unsqueeze(0) + + return (cropped_tensor,) \ No newline at end of file diff --git a/load_images.py b/load_images.py new file mode 100644 index 0000000..9dc74d6 --- /dev/null +++ b/load_images.py @@ -0,0 +1,69 @@ +import os +import hashlib +import numpy as np +import torch +from PIL import Image +import folder_paths + +class LoadImagesOriginalSize: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING", {"default": "", "placeholder": "Path to directory of images"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + OUTPUT_IS_LIST = (True,) + FUNCTION = "load_images" + CATEGORY = "image/loading" + + def load_images(self, directory: str): + if not os.path.isdir(directory): + raise FileNotFoundError(f"Directory '{directory}' cannot be found.") + + valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.webp') + image_files = [f for f in os.listdir(directory) + if f.lower().endswith(valid_extensions)] + + if len(image_files) == 0: + raise FileNotFoundError(f"No valid image files found in directory '{directory}'.") + + image_files.sort() + + images_list = [] + for filename in image_files: + filepath = os.path.join(directory, filename) + + img = Image.open(filepath) + if img.mode != 'RGB': + img = img.convert('RGB') + + img_np = np.array(img, dtype=np.float32) / 255.0 + img_tensor = torch.from_numpy(img_np).unsqueeze(0) + images_list.append(img_tensor) + + return (images_list,) + + @classmethod + def IS_CHANGED(cls, directory: str): + if not os.path.isdir(directory): + return False + + m = hashlib.sha256() + valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.webp') + for filename in sorted(os.listdir(directory)): + if filename.lower().endswith(valid_extensions): + filepath = os.path.join(directory, filename) + m.update(str(os.path.getmtime(filepath)).encode()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(cls, directory: str): + if not directory.strip(): + return "Directory path cannot be empty" + + if not os.path.isdir(directory): + return f"Directory '{directory}' cannot be found." + return True \ No newline at end of file diff --git a/nearest_upscale.py b/nearest_upscale.py new file mode 100644 index 0000000..ad69de1 --- /dev/null +++ b/nearest_upscale.py @@ -0,0 +1,71 @@ +import torch +import numpy as np + +class NearestUpscaleNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "scale_factor": ("INT", { + "default": 2, + "min": 1, + "max": 8, + "step": 1 + }), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "upscale" + CATEGORY = "image/upscaling" + + def upscale(self, image, scale_factor): + # Debug inicial do tensor de entrada + print(f"Tipo de entrada: {type(image)}") + print(f"Shape do tensor de entrada: {image.shape}") + + # Converter tensor para numpy array + if isinstance(image, torch.Tensor): + image_np = image[0].cpu().numpy() + else: + image_np = np.array(image) + + # Debug após conversão para numpy + print(f"Shape do numpy array: {image_np.shape}") + height, width = image_np.shape[:2] + print(f"Altura: {height}, Largura: {width}") + print(f"Fator de escala: {scale_factor}") + + # Calcular novas dimensões + new_height = height * scale_factor + new_width = width * scale_factor + print(f"Nova altura calculada: {new_height}") + print(f"Nova largura calculada: {new_width}") + + # Criar array de saída com as dimensões exatas + upscaled = np.zeros((new_height, new_width, 3), dtype=np.float32) + print(f"Shape do array de saída: {upscaled.shape}") + + # Upscaling usando repeat nativo do numpy + # Primeiro expandimos na direção vertical + temp = np.repeat(image_np, scale_factor, axis=0) + # Depois na horizontal + upscaled = np.repeat(temp, scale_factor, axis=1) + + print(f"Shape final antes do tensor: {upscaled.shape}") + + # Converter para tensor mantendo as dimensões exatas + upscaled_tensor = torch.from_numpy(upscaled).unsqueeze(0) + print(f"Shape final do tensor: {upscaled_tensor.shape}") + + return (upscaled_tensor,) + +# Para registrar o nó +NODE_CLASS_MAPPINGS = { + "NearestUpscale": NearestUpscaleNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NearestUpscale": "Nearest Neighbor Upscale" +} \ No newline at end of file diff --git a/pixel_normalizer.py b/pixel_normalizer.py new file mode 100644 index 0000000..7987842 --- /dev/null +++ b/pixel_normalizer.py @@ -0,0 +1,170 @@ +import numpy as np +import torch +from PIL import Image +import cv2 +from sklearn.cluster import KMeans + +class PixelArtNormalizerNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "block_size": ("INT", { + "default": 4, + "min": 0, + "max": 8, + "step": 1, + "description": "0 for auto-detection" + }), + "n_colors": ("INT", { + "default": 32, + "min": 0, + "max": 256, + "step": 1, + "description": "0 for auto-detection" + }) + } + } + + RETURN_TYPES = ("IMAGE", "INT", "IMAGE") # imagem normal, tamanho do bloco, imagem downscaled + RETURN_NAMES = ("normalized", "block_size", "downscaled") + FUNCTION = "normalize_pixel_art" + CATEGORY = "image/processing" + + def detect_grid(self, image): + """Detecta o tamanho aproximado dos pixels na grade.""" + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, 50, 150) + + # Detecta linhas + lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=50, + minLineLength=20, maxLineGap=5) + + if lines is None: + return 2 # valor padrão se não detectar + + # Calcula distâncias entre linhas paralelas + distances = [] + for i in range(len(lines)): + x1, y1, x2, y2 = lines[i][0] + for j in range(i + 1, len(lines)): + x3, y3, x4, y4 = lines[j][0] + + angle1 = np.arctan2(y2 - y1, x2 - x1) + angle2 = np.arctan2(y4 - y3, x4 - x3) + if abs(angle1 - angle2) < 0.1: + dist = abs((y4 - y3) * x1 - (x4 - x3) * y1 + x4 * y3 - y4 * x3) / \ + np.sqrt((y4 - y3)**2 + (x4 - x3)**2) + if dist > 2: + distances.append(dist) + + if not distances: + return 2 + + grid_size = int(np.median(distances)) + return max(2, min(grid_size, 8)) # limita entre 2 e 8 pixels + + def normalize_to_grid(self, image, grid_size): + h, w = image.shape[:2] + + # Ajusta dimensões para serem múltiplos do grid_size + new_h = ((h + grid_size - 1) // grid_size) * grid_size + new_w = ((w + grid_size - 1) // grid_size) * grid_size + + # Cria nova imagem com padding se necessário + normalized = np.zeros((new_h, new_w, 3), dtype=np.uint8) + normalized[:h, :w] = image + + # Para cada célula da grade + for y in range(0, new_h, grid_size): + for x in range(0, new_w, grid_size): + # Limita as coordenadas aos limites da imagem original + y_end = min(y + grid_size, h) + x_end = min(x + grid_size, w) + + # Pega o bloco atual + block = normalized[y:y_end, x:x_end] + + if block.size > 0: + # Encontra a cor mais frequente no bloco + block_reshaped = block.reshape(-1, 3) + unique_colors, counts = np.unique(block_reshaped, axis=0, return_counts=True) + dominant_color = unique_colors[counts.argmax()] + + # Preenche o bloco com a cor dominante + normalized[y:y_end, x:x_end] = dominant_color + + return normalized[:h, :w] + + def quantize_colors(self, image, n_colors): + h, w = image.shape[:2] + pixels = image.reshape(-1, 3) + + kmeans = KMeans(n_clusters=n_colors, random_state=42) + labels = kmeans.fit_predict(pixels) + palette = kmeans.cluster_centers_.astype(np.uint8) + + quantized = palette[labels].reshape(h, w, 3) + return quantized + + def normalize_pixel_art(self, image, block_size, n_colors): + # Converter tensor para numpy array + if isinstance(image, torch.Tensor): + if image.dim() == 4: + image_np = image[0].cpu().numpy() + else: + image_np = image.cpu().numpy() + else: + image_np = np.array(image) + + # Converter para uint8 se estiver normalizado entre 0-1 + if image_np.max() <= 1.0: + image_np = (image_np * 255).astype(np.uint8) + + # Usar detecção automática se block_size ou n_colors forem 0 + if n_colors <= 0: + n_colors = min(32, max(8, int(np.sqrt(image_np.shape[0] * image_np.shape[1] / 100)))) + print(f"Número de cores detectado automaticamente: {n_colors}") + + # Quantizar cores + quantized = self.quantize_colors(image_np, n_colors) + + # Detectar tamanho do grid se block_size for 0 + detected_block_size = self.detect_grid(quantized) if block_size <= 0 else block_size + print(f"Tamanho do grid: {detected_block_size}px") + + # Normalizar para a grade + normalized = self.normalize_to_grid(quantized, detected_block_size) + + # Criar versão downscaled + h, w = normalized.shape[:2] + new_h = h // detected_block_size + new_w = w // detected_block_size + + # Usar área de cada bloco para determinar a cor do pixel correspondente + downscaled = np.zeros((new_h, new_w, 3), dtype=np.uint8) + for y in range(new_h): + for x in range(new_w): + block = normalized[y*detected_block_size:(y+1)*detected_block_size, + x*detected_block_size:(x+1)*detected_block_size] + downscaled[y, x] = block[0, 0] # Como o bloco já está normalizado, podemos pegar qualquer pixel + + # Converter ambas as imagens para float32 normalizado + normalized = normalized.astype(np.float32) / 255.0 + downscaled = downscaled.astype(np.float32) / 255.0 + + # Converter para tensores + normalized_tensor = torch.from_numpy(normalized).unsqueeze(0) + downscaled_tensor = torch.from_numpy(downscaled).unsqueeze(0) + + return (normalized_tensor, detected_block_size, downscaled_tensor) + +# Registrar o nó +NODE_CLASS_MAPPINGS = { + "PixelArtNormalizer": PixelArtNormalizerNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PixelArtNormalizer": "Pixel Art Normalizer" +} \ No newline at end of file diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..2fee449 --- /dev/null +++ b/readme.md @@ -0,0 +1,116 @@ +# Image Processing Suite for ComfyUI + +A collection of specialized image processing nodes for ComfyUI, focused on dataset preparation and pixel art manipulation. + +## Installation + +1. Create a `custom_nodes` directory in your ComfyUI installation if it doesn't exist +2. Clone this repository inside the `custom_nodes` directory: +```bash +cd custom_nodes +git clone [repository_url] image_processing +``` +3. Restart ComfyUI + +## Nodes + +### Load Images (Original Size) +Loads all images from a directory while preserving their original dimensions. + +**Inputs:** +- `directory`: Path to the directory containing images + +**Outputs:** +- List of images in their original sizes + +**Features:** +- Preserves original image dimensions +- Compatible with ComfyUI's RebatchImages node +- Supports common image formats (png, jpg, jpeg, bmp, webp) + +### Custom Crop +Crops images with specific positioning options. + +**Inputs:** +- `image`: Input image +- `crop_width`: Width of crop area +- `crop_height`: Height of crop area +- `crop_mode`: Cropping position ("center", "left", "right", "top", "bottom") + +**Outputs:** +- Cropped image + +### Smart Resize +Resizes images to a target size while intelligently filling missing areas. + +**Inputs:** +- `image`: Input image +- `target_size`: Desired size +- `border_sample_size`: Pixels to sample for border color +- `color_method`: Method to determine fill color ("mean" or "mode") + +**Outputs:** +- Resized image with intelligent border filling + +### Nearest Neighbor Upscale +Performs upscaling using nearest neighbor interpolation, perfect for pixel art. + +**Inputs:** +- `image`: Input image +- `scale_factor`: Multiplication factor for upscaling (1-8) + +**Outputs:** +- Upscaled image without interpolation artifacts + +### Pixel Art Normalizer +Normalizes images into pixel art style with consistent grid sizes. + +**Inputs:** +- `image`: Input image +- `block_size`: Size of pixel blocks (0 for auto-detection) +- `n_colors`: Number of colors in output (0 for auto-detection) + +**Outputs:** +- `normalized`: Normalized pixel art image at original size +- `block_size`: Detected/used block size +- `downscaled`: 1:1 pixel art version (downscaled by block size) + +**Features:** +- Automatic grid size detection +- Color quantization +- Outputs both full-size and true 1:1 pixel art versions + +## Usage Examples + +### Basic Image Loading and Batching +``` +LoadImagesOriginal -> RebatchImages -> [Further Processing] +``` + +### Pixel Art Creation Pipeline +``` +LoadImagesOriginal -> PixelArtNormalizer -> NearestUpscale +``` + +### Dataset Preparation +``` +LoadImagesOriginal -> CustomCrop -> SmartResize -> [Training] +``` + +## Dependencies +- NumPy +- OpenCV (cv2) +- scikit-learn +- PIL +- PyTorch (provided by ComfyUI) + +## Notes +- All nodes maintain compatibility with ComfyUI's native nodes +- Images are handled in RGB format +- All operations preserve proper normalization (0-1 range) + +## Contributing +Feel free to open issues or submit pull requests for improvements. + +## License +[Your chosen license] diff --git a/smart_resize.py b/smart_resize.py new file mode 100644 index 0000000..ff14726 --- /dev/null +++ b/smart_resize.py @@ -0,0 +1,136 @@ +import torch +import numpy as np +from PIL import Image +from scipy import stats +from collections import Counter + +class SmartResizeNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "target_size": ("INT", { + "default": 1024, + "min": 64, + "max": 8192, + "step": 64 + }), + "border_sample_size": ("INT", { + "default": 5, + "min": 1, + "max": 50, + "step": 1, + "description": "Pixels to sample for border color" + }), + "color_method": (["mean", "mode"],), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "process_image" + CATEGORY = "image/processing" + + def get_border_color_mean(self, img_np, edge, sample_size): + h, w = img_np.shape[:2] + + if edge in ['top', 'bottom']: + rows = slice(0, sample_size) if edge == 'top' else slice(h-sample_size, h) + return np.mean(img_np[rows, :], axis=(0, 1)) + else: + cols = slice(0, sample_size) if edge == 'left' else slice(w-sample_size, w) + return np.mean(img_np[:, cols], axis=(0, 1)) + + def get_border_color_mode(self, img_np, edge, sample_size): + h, w = img_np.shape[:2] + + # Extrair a região da borda + if edge in ['top', 'bottom']: + rows = slice(0, sample_size) if edge == 'top' else slice(h-sample_size, h) + border_pixels = img_np[rows, :] + else: + cols = slice(0, sample_size) if edge == 'left' else slice(w-sample_size, w) + border_pixels = img_np[:, cols] + + # Reshapear para ter uma lista de pixels + pixels = border_pixels.reshape(-1, 3) + + # Arredondar para reduzir ruído e facilitar encontrar cores iguais + pixels = np.round(pixels * 255) / 255 + + # Usar Counter para encontrar a cor mais frequente + pixel_tuples = [tuple(p) for p in pixels] + most_common_color = Counter(pixel_tuples).most_common(1)[0][0] + + # Converter de volta para array numpy + return np.array(most_common_color, dtype=np.float32) + + def get_border_color(self, img_np, edge, sample_size, method='mean'): + if method == 'mean': + return self.get_border_color_mean(img_np, edge, sample_size) + else: # method == 'mode' + return self.get_border_color_mode(img_np, edge, sample_size) + + def process_image(self, image, target_size, border_sample_size, color_method): + # Converter tensor para numpy + if isinstance(image, torch.Tensor): + image_np = image[0].cpu().numpy() + else: + image_np = np.array(image) + + # Garantir que estamos trabalhando com valores entre 0 e 1 + if image_np.max() > 1.0: + image_np = image_np.astype(np.float32) / 255.0 + + height, width = image_np.shape[:2] + current_image = image_np.copy() + + # Processar largura + if width > target_size: + # Cortar o excesso da largura + start_x = (width - target_size) // 2 + current_image = current_image[:, start_x:start_x + target_size] + elif width < target_size: + # Criar nova imagem com a largura correta + temp_image = np.zeros((height, target_size, 3), dtype=np.float32) + + # Calcular padding + padding_left = (target_size - width) // 2 + + # Pegar cores das bordas + left_color = self.get_border_color(current_image, 'left', border_sample_size, color_method) + right_color = self.get_border_color(current_image, 'right', border_sample_size, color_method) + + # Preencher as bordas + temp_image[:, :padding_left] = left_color + temp_image[:, padding_left:padding_left + width] = current_image + temp_image[:, padding_left + width:] = right_color + + current_image = temp_image + + # Processar altura + current_height = current_image.shape[0] + if current_height > target_size: + # Cortar o excesso da altura + start_y = (current_height - target_size) // 2 + current_image = current_image[start_y:start_y + target_size] + elif current_height < target_size: + # Criar nova imagem com a altura correta + temp_image = np.zeros((target_size, target_size, 3), dtype=np.float32) + + # Calcular padding + padding_top = (target_size - current_height) // 2 + + # Pegar cores das bordas + top_color = self.get_border_color(current_image, 'top', border_sample_size, color_method) + bottom_color = self.get_border_color(current_image, 'bottom', border_sample_size, color_method) + + # Preencher as bordas + temp_image[:padding_top] = top_color + temp_image[padding_top:padding_top + current_height] = current_image + temp_image[padding_top + current_height:] = bottom_color + + current_image = temp_image + + # Converter de volta para tensor + return (torch.from_numpy(current_image).unsqueeze(0),) \ No newline at end of file