diff --git a/README.md b/README.md index ce6fd5a..3c10fd3 100644 --- a/README.md +++ b/README.md @@ -136,7 +136,6 @@ For more information about how IRC uses colors refer to these resources: - torch - numpy - scikit-learn -- pillow ## Credits diff --git a/requirements.txt b/requirements.txt index 2268d65..c7b1ce7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,3 @@ torch>=1.9.0 numpy>=1.21.0 -scikit-learn>=1.0.0 -pillow>=8.0.0 \ No newline at end of file +scikit-learn>=1.0.0 \ No newline at end of file diff --git a/src/nodes.py b/src/nodes.py index 6c3f11d..93ee1a0 100644 --- a/src/nodes.py +++ b/src/nodes.py @@ -1,7 +1,6 @@ import torch import numpy as np from sklearn.cluster import KMeans -from PIL import Image from datetime import datetime class IRCArtConverter: @@ -74,9 +73,7 @@ class IRCArtConverter: else: # Fallback to weighted manhattan - weights = np.array([0.7, 1.0, 0.5]) - diff = np.abs(detected_color - target_colors) - return np.sum(diff * weights, axis=1) + return self.calculate_color_distance(detected_color, target_colors, "weighted_manhattan") def apply_color_transfer(self, image, use_16_colors=False, distance_method="weighted_manhattan"): """Apply color transfer using KMeans clustering and configurable distance matching.""" @@ -93,7 +90,7 @@ class IRCArtConverter: detected_colors = clustering_model.cluster_centers_.astype(int) - # Match detected colors to IRC palette using Manhattan distance + # Match detected colors to IRC palette using configured distance method target_colors_array = np.array(target_colors) closest_colors = [] @@ -152,25 +149,40 @@ class IRCArtConverter: # Get representative color for this block block_color_rounded = tuple(np.round(block_color).astype(int)) - # Find the exact IRC color index using exact matching first + # Select palette based on mode + palette = self.IRC_COLORS[:16] if use_16_colors else self.IRC_COLORS + # Try exact match first (since colors should already be in IRC palette) try: - if use_16_colors: - color_idx = self.IRC_COLORS[:16].index(block_color_rounded) - else: - color_idx = self.IRC_COLORS.index(block_color_rounded) + return palette.index(block_color_rounded) except ValueError: # Fallback to distance matching if exact match fails - block_color_array = np.array(block_color) - if use_16_colors: - target_colors_array = np.array(self.IRC_COLORS[:16]) - else: - target_colors_array = np.array(self.IRC_COLORS) - - distances = self.calculate_color_distance(block_color_array, target_colors_array, distance_method) - color_idx = np.argmin(distances) + target_colors_array = np.array(palette) + distances = self.calculate_color_distance(np.array(block_color), target_colors_array, distance_method) + return np.argmin(distances) + + def _format_color_code(self, fg_color, bg_color=None): + """Format IRC color code with proper padding.""" + if bg_color is None: + bg_color = fg_color - return color_idx + fg_fmt = f"{fg_color:02d}" if fg_color >= 10 else str(fg_color) + bg_fmt = f"{bg_color:02d}" if bg_color >= 10 else str(bg_color) + return f"\x03{fg_fmt},{bg_fmt}" + + def _count_consecutive(self, colors, index, is_tuple_check=False): + """Count consecutive identical entries starting from index.""" + count = 1 + current = colors[index] + + while index + count < len(colors): + next_entry = colors[index + count] + # Check type consistency and value equality + if (is_tuple_check and not isinstance(next_entry, tuple)) or next_entry != current: + break + count += 1 + + return count def format_irc_line(self, line_colors): """Format a line of IRC color codes with compression.""" @@ -179,59 +191,27 @@ class IRCArtConverter: line_text = "" i = 0 + while i < len(line_colors): color_entry = line_colors[i] - # Check if this is half-block mode (tuple) or full-block mode (int) if isinstance(color_entry, tuple): # Half-block mode: (top_color, bottom_color) top_color, bottom_color = color_entry if top_color == bottom_color: - # Same color - use space with background color, check for compression - count = 1 - # Count consecutive identical same-color tuples - while (i + count < len(line_colors) and - isinstance(line_colors[i + count], tuple) and - line_colors[i + count] == color_entry): - count += 1 - - # Format color code once - if top_color < 10: - color_format = f"\x03{top_color},{top_color}" - else: - color_format = f"\x03{top_color:02d},{top_color:02d}" - - line_text += color_format + (" " * count) + # Same color - compress consecutive identical blocks + count = self._count_consecutive(line_colors, i, is_tuple_check=True) + line_text += self._format_color_code(top_color) + (" " * count) i += count else: - # Different colors - use half block character (no compression for mixed colors) - if top_color < 10 and bottom_color < 10: - color_format = f"\x03{top_color},{bottom_color}" - elif top_color < 10: - color_format = f"\x03{top_color},{bottom_color:02d}" - elif bottom_color < 10: - color_format = f"\x03{top_color:02d},{bottom_color}" - else: - color_format = f"\x03{top_color:02d},{bottom_color:02d}" - line_text += color_format + "▀" + # Different colors - use half block character + line_text += self._format_color_code(top_color, bottom_color) + "▀" i += 1 else: # Full-block mode: single color - color_code = color_entry - count = 1 - - # Count consecutive identical colors - while i + count < len(line_colors) and not isinstance(line_colors[i + count], tuple) and line_colors[i + count] == color_code: - count += 1 - - # Format color code - if color_code < 10: - color_format = f"\x03{color_code},{color_code}" - else: - color_format = f"\x03{color_code:02d},{color_code:02d}" - - line_text += color_format + (" " * count) + count = self._count_consecutive(line_colors, i) + line_text += self._format_color_code(color_entry) + (" " * count) i += count return line_text @@ -255,6 +235,70 @@ class IRCArtConverter: block_color = self.get_block_color(block) return self.find_irc_color_index(block_color, use_16_colors, distance_method), start_y, end_y + def _apply_color_to_output(self, output, color_idx, start_y, end_y, start_x, end_x): + """Apply IRC color to output image region.""" + irc_color = self.IRC_COLORS[color_idx] + irc_color_tensor = torch.tensor(irc_color, dtype=torch.float32) / 255.0 + output[start_y:end_y, start_x:end_x, :] = irc_color_tensor + + def _process_half_block_row(self, color_transferred_img, y, blocks_x, blocks_y, output, params): + """Process a row of half-blocks.""" + line_colors = [] + block_width, block_height, actual_block_height = params['block_width'], params['block_height'], params['actual_block_height'] + width, height = params['width'], params['height'] + use_16_colors, distance_method = params['use_16_colors'], params['distance_method'] + + for x in range(blocks_x): + # Process first half-block + color_idx, start_y, end_y = self.process_block( + color_transferred_img, x, y, block_width, block_height, + actual_block_height, width, height, use_16_colors, distance_method + ) + + # Apply color to image + x_start = x * block_width + x_end = min((x + 1) * block_width, width) + self._apply_color_to_output(output, color_idx, start_y, end_y, x_start, x_end) + + # Process second half-block if it exists + if y + 1 < blocks_y: + color_idx2, start_y2, end_y2 = self.process_block( + color_transferred_img, x, y + 1, block_width, block_height, + actual_block_height, width, height, use_16_colors, distance_method + ) + + # Apply color to second block + self._apply_color_to_output(output, color_idx2, start_y2, end_y2, x_start, x_end) + + # Store both colors as tuple (top, bottom) + line_colors.append((color_idx, color_idx2)) + else: + # Last row, only one half-block - use same color for both + line_colors.append((color_idx, color_idx)) + + return line_colors + + def _process_full_block_row(self, color_transferred_img, y, blocks_x, output, params): + """Process a row of full blocks.""" + line_colors = [] + block_width, block_height, actual_block_height = params['block_width'], params['block_height'], params['actual_block_height'] + width, height = params['width'], params['height'] + use_16_colors, distance_method = params['use_16_colors'], params['distance_method'] + + for x in range(blocks_x): + color_idx, start_y, end_y = self.process_block( + color_transferred_img, x, y, block_width, block_height, + actual_block_height, width, height, use_16_colors, distance_method + ) + + # Apply color to image and store for text + x_start = x * block_width + x_end = min((x + 1) * block_width, width) + self._apply_color_to_output(output, color_idx, start_y, end_y, x_start, x_end) + line_colors.append(color_idx) + + return line_colors + def convert_to_irc_art(self, image, block_width, block_height, half_block_mode, use_16_colors, distance_method): # Input validation batch_size, height, width, channels = image.shape @@ -282,60 +326,27 @@ class IRCArtConverter: output = torch.zeros_like(img) text_lines = [] + # Bundle parameters for cleaner function calls + params = { + 'block_width': block_width, + 'block_height': block_height, + 'actual_block_height': actual_block_height, + 'width': width, + 'height': height, + 'use_16_colors': use_16_colors, + 'distance_method': distance_method + } + # Process blocks if half_block_mode: # Half-block mode: process in pairs for y in range(0, blocks_y, 2): - line_colors = [] - - for x in range(blocks_x): - # Process first half-block - color_idx, start_y, end_y = self.process_block( - color_transferred_img, x, y, block_width, block_height, - actual_block_height, width, height, use_16_colors, distance_method - ) - - # Apply color to image - irc_color = self.IRC_COLORS[color_idx] - irc_color_tensor = torch.tensor(irc_color, dtype=torch.float32) / 255.0 - output[start_y:end_y, x * block_width:min((x + 1) * block_width, width), :] = irc_color_tensor - - # Process second half-block if it exists - if y + 1 < blocks_y: - color_idx2, start_y2, end_y2 = self.process_block( - color_transferred_img, x, y + 1, block_width, block_height, - actual_block_height, width, height, use_16_colors, distance_method - ) - - # Apply color to second block - irc_color2 = self.IRC_COLORS[color_idx2] - irc_color_tensor2 = torch.tensor(irc_color2, dtype=torch.float32) / 255.0 - output[start_y2:end_y2, x * block_width:min((x + 1) * block_width, width), :] = irc_color_tensor2 - - # Store both colors as tuple (top, bottom) - line_colors.append((color_idx, color_idx2)) - else: - # Last row, only one half-block - use same color for both - line_colors.append((color_idx, color_idx)) - + line_colors = self._process_half_block_row(color_transferred_img, y, blocks_x, blocks_y, output, params) text_lines.append(self.format_irc_line(line_colors)) else: # Full block mode for y in range(blocks_y): - line_colors = [] - - for x in range(blocks_x): - color_idx, start_y, end_y = self.process_block( - color_transferred_img, x, y, block_width, block_height, - actual_block_height, width, height, use_16_colors, distance_method - ) - - # Apply color to image and store for text - irc_color = self.IRC_COLORS[color_idx] - irc_color_tensor = torch.tensor(irc_color, dtype=torch.float32) / 255.0 - output[start_y:end_y, x * block_width:min((x + 1) * block_width, width), :] = irc_color_tensor - line_colors.append(color_idx) # Use same index as image - + line_colors = self._process_full_block_row(color_transferred_img, y, blocks_x, output, params) text_lines.append(self.format_irc_line(line_colors)) # Join all lines with newlines