From 869e3f6f881f40fc642f1f903bd158c4cd59bb3c Mon Sep 17 00:00:00 2001 From: Hugh Bord Date: Wed, 27 Aug 2025 10:00:17 +0800 Subject: [PATCH] feat: 16 color support and optimised text output --- src/nodes.py | 84 ++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 61 insertions(+), 23 deletions(-) diff --git a/src/nodes.py b/src/nodes.py index fb3a7e0..f68ada3 100644 --- a/src/nodes.py +++ b/src/nodes.py @@ -2,6 +2,7 @@ import torch import numpy as np from sklearn.cluster import KMeans from PIL import Image +from datetime import datetime class IRCArtConverter: CATEGORY = "image/conversion" @@ -14,6 +15,7 @@ class IRCArtConverter: "block_width": ("INT", {"default": 8, "min": 1, "max": 32}), "block_height": ("INT", {"default": 15, "min": 1, "max": 32}), "half_block_mode": ("BOOLEAN", {"default": False}), + "use_16_colors": ("BOOLEAN", {"default": False}), "color_method": (["dominant", "average", "median"],) } } @@ -52,7 +54,7 @@ class IRCArtConverter: """Manhattan distance function matching the color transfer node""" return np.sum(np.abs(detected_color - target_colors), axis=1) - def apply_color_transfer(self, image): + def apply_color_transfer(self, image, use_16_colors=False): """Apply color transfer using KMeans clustering and Manhattan distance matching""" # Convert image to numpy array img_array = (image * 255.0).cpu().numpy().astype(np.uint8) @@ -60,14 +62,20 @@ class IRCArtConverter: # Reshape for clustering img_flat = img_array.reshape((-1, 3)) + # Select color palette based on mode + if use_16_colors: + target_colors = self.IRC_COLORS[:16] # Use only first 16 colors + else: + target_colors = self.IRC_COLORS # Use all 99 colors + # KMeans clustering - clustering_model = KMeans(n_clusters=len(self.IRC_COLORS), n_init="auto", random_state=42) + clustering_model = KMeans(n_clusters=len(target_colors), n_init="auto", random_state=42) clustering_model.fit(img_flat) detected_colors = clustering_model.cluster_centers_.astype(int) # Match detected colors to IRC palette using Manhattan distance - target_colors_array = np.array(self.IRC_COLORS) + target_colors_array = np.array(target_colors) closest_colors = [] for color in detected_colors: @@ -136,7 +144,7 @@ class IRCArtConverter: # Fallback to average if k-means fails return tuple(block.mean(dim=(0, 1)).cpu().numpy() * 255) - def convert_to_irc_art(self, image, block_width, block_height, half_block_mode, color_method): + def convert_to_irc_art(self, image, block_width, block_height, half_block_mode, use_16_colors, color_method): # Input validation batch_size, height, width, channels = image.shape @@ -147,7 +155,7 @@ class IRCArtConverter: img = image[0] # Shape: (height, width, 3) # Step 1: Apply color transfer to map image to IRC palette - color_transferred_img = self.apply_color_transfer(img) + color_transferred_img = self.apply_color_transfer(img, use_16_colors) # Adjust block height for half block mode actual_block_height = block_height / 2 if half_block_mode else block_height @@ -188,11 +196,23 @@ class IRCArtConverter: # Find the exact IRC color index since colors are already mapped block_color_array = np.array(block_color) - target_colors_array = np.array(self.IRC_COLORS) + + # Use the same color palette as color transfer + if use_16_colors: + target_colors_array = np.array(self.IRC_COLORS[:16]) + else: + target_colors_array = np.array(self.IRC_COLORS) + distances = self.manhattan_distance(block_color_array, target_colors_array) color_idx = np.argmin(distances) - irc_color = self.IRC_COLORS[color_idx] - irc_code = self.get_irc_color_code(color_idx) + + # Get the actual color and code from the full palette + if use_16_colors: + irc_color = self.IRC_COLORS[color_idx] # color_idx is 0-15 + irc_code = self.get_irc_color_code(color_idx) # codes 1-16 + else: + irc_color = self.IRC_COLORS[color_idx] # color_idx is 0-98 + irc_code = self.get_irc_color_code(color_idx) # codes 1-99 # Store color code for text output line_colors.append(irc_code) @@ -201,16 +221,35 @@ class IRCArtConverter: irc_color_tensor = torch.tensor(irc_color, dtype=torch.float32) / 255.0 output[start_y:end_y, start_x:end_x, :] = irc_color_tensor - # Build IRC text line with color codes - # Format: \x03{fg},{bg}{char} where fg=bg for solid blocks + # Build IRC text line with optimized color codes and compression line_text = "" - for color_code in line_colors: - # Use the same color for foreground and background, with space character - line_text += f"\x03{color_code:02d},{color_code:02d} " + if not line_colors: + text_lines.append("") + continue + + # Compress consecutive identical colors + i = 0 + while i < len(line_colors): + color_code = line_colors[i] + count = 1 + + # Count consecutive identical colors + while i + count < len(line_colors) and line_colors[i + count] == color_code: + count += 1 + + # Format color code (single digits for 1-9, double for 10+) + if color_code <= 9: + color_format = f"\x03{color_code},{color_code}" + else: + color_format = f"\x03{color_code:02d},{color_code:02d}" + + # Add spaces based on count + line_text += color_format + (" " * count) + i += count text_lines.append(line_text) - # Join all lines with newlines + # Join all lines with newlines (no header timestamp) irc_text = "\n".join(text_lines) # Return as batch of 1 @@ -238,7 +277,13 @@ class IRCTextSaver: def save_irc_text(self, irc_text, filename): """Save IRC text to a file""" import os - from datetime import datetime + + # Generate current timestamp + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + + # Replace {TIMESTAMP} placeholder in filename if present + if "{TIMESTAMP}" in filename: + filename = filename.replace("{TIMESTAMP}", timestamp) # Ensure filename has .txt extension if not filename.endswith('.txt'): @@ -248,16 +293,9 @@ class IRCTextSaver: output_dir = "output/irc_art" os.makedirs(output_dir, exist_ok=True) - # Create full file path with timestamp if file exists - base_name = filename.replace('.txt', '') + # Create full file path file_path = os.path.join(output_dir, filename) - # If file exists, add timestamp - if os.path.exists(file_path): - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"{base_name}_{timestamp}.txt" - file_path = os.path.join(output_dir, filename) - # Write the IRC text to file try: with open(file_path, 'w', encoding='utf-8') as f: