feat: work on refactor

This commit is contained in:
Hugh Bord
2025-08-27 11:19:40 +08:00
parent 26092557a4
commit 06883d8a5f
2 changed files with 141 additions and 126 deletions
+5
View File
@@ -107,6 +107,11 @@ The node uses a two-stage approach:
The node uses an extended 99-color IRC palette built in.
For more information about how IRC uses colors refer to these resources:
- [https://modern.ircdocs.horse/formatting.html#color](https://modern.ircdocs.horse/formatting.html#color)
- [https://anti.teamidiot.de/static/nei/*/extended_mirc_color_proposal.html](https://anti.teamidiot.de/static/nei/*/extended_mirc_color_proposal.html)
## Requirements
- ComfyUI
+136 -126
View File
@@ -7,6 +7,11 @@ from datetime import datetime
class IRCArtConverter:
CATEGORY = "image/conversion"
# 99-color IRC palette (class constant)
IRC_COLORS = [
(255, 255, 255), (0, 0, 0), (0, 0, 127), (0, 147, 0), (255, 0, 0), (127, 0, 0), (156, 0, 156), (252, 127, 0), (255, 255, 0), (0, 252, 0), (0, 147, 147), (0, 255, 255), (0, 0, 252), (255, 0, 255), (127, 127, 127), (210, 210, 210), (71, 0, 0), (71, 33, 0), (71, 71, 0), (50, 71, 0), (0, 71, 0), (0, 71, 44), (0, 71, 71), (0, 39, 71), (0, 0, 71), (46, 0, 71), (71, 0, 71), (71, 0, 42), (116, 0, 0), (116, 58, 0), (116, 116, 0), (81, 116, 0), (0, 116, 0), (0, 116, 73), (0, 116, 116), (0, 64, 116), (0, 0, 116), (75, 0, 116), (116, 0, 116), (116, 0, 69), (181, 0, 0), (181, 99, 0), (181, 181, 0), (125, 181, 0), (0, 181, 0), (0, 181, 113), (0, 181, 181), (0, 99, 181), (0, 0, 181), (117, 0, 181), (181, 0, 181), (181, 0, 107), (255, 0, 0), (255, 140, 0), (255, 255, 0), (178, 255, 0), (0, 255, 0), (0, 255, 160), (0, 255, 255), (0, 140, 255), (0, 0, 255), (165, 0, 255), (255, 0, 255), (255, 0, 152), (255, 89, 89), (255, 180, 89), (255, 255, 113), (207, 255, 96), (111, 255, 111), (101, 255, 201), (109, 255, 255), (89, 180, 255), (89, 89, 255), (196, 89, 255), (255, 102, 255), (255, 89, 188), (255, 156, 156), (255, 211, 156), (255, 255, 156), (226, 255, 156), (156, 255, 156), (156, 255, 219), (156, 255, 255), (156, 211, 255), (156, 156, 255), (220, 156, 255), (255, 156, 255), (255, 148, 211), (0, 0, 0), (19, 19, 19), (40, 40, 40), (54, 54, 54), (77, 77, 77), (101, 101, 101), (129, 129, 129), (159, 159, 159), (188, 188, 188), (226, 226, 226), (255, 255, 255)
]
@classmethod
def INPUT_TYPES(s):
return {
@@ -24,49 +29,18 @@ class IRCArtConverter:
RETURN_NAMES = ("processed_image", "irc_text")
FUNCTION = "convert_to_irc_art"
def __init__(self):
# 99-color IRC palette
self.IRC_COLORS = [
(255, 255, 255), (0, 0, 0), (0, 0, 127), (0, 147, 0), (255, 0, 0),
(127, 0, 0), (156, 0, 156), (252, 127, 0), (255, 255, 0), (0, 252, 0),
(0, 147, 147), (0, 255, 255), (0, 0, 252), (255, 0, 255), (127, 127, 127),
(210, 210, 210), (71, 0, 0), (71, 33, 0), (71, 71, 0), (50, 71, 0),
(0, 71, 0), (0, 71, 44), (0, 71, 71), (0, 39, 71), (0, 0, 71),
(46, 0, 71), (71, 0, 71), (71, 0, 42), (116, 0, 0), (116, 58, 0),
(116, 116, 0), (81, 116, 0), (0, 116, 0), (0, 116, 73), (0, 116, 116),
(0, 64, 116), (0, 0, 116), (75, 0, 116), (116, 0, 116), (116, 0, 69),
(181, 0, 0), (181, 99, 0), (181, 181, 0), (125, 181, 0), (0, 181, 0),
(0, 181, 113), (0, 181, 181), (0, 99, 181), (0, 0, 181), (117, 0, 181),
(181, 0, 181), (181, 0, 107), (255, 0, 0), (255, 140, 0), (255, 255, 0),
(178, 255, 0), (0, 255, 0), (0, 255, 160), (0, 255, 255), (0, 140, 255),
(0, 0, 255), (165, 0, 255), (255, 0, 255), (255, 0, 152), (255, 89, 89),
(255, 180, 89), (255, 255, 113), (207, 255, 96), (111, 255, 111), (101, 255, 201),
(109, 255, 255), (89, 180, 255), (89, 89, 255), (196, 89, 255), (255, 102, 255),
(255, 89, 188), (255, 156, 156), (255, 211, 156), (255, 255, 156), (226, 255, 156),
(156, 255, 156), (156, 255, 219), (156, 255, 255), (156, 211, 255), (156, 156, 255),
(220, 156, 255), (255, 156, 255), (255, 148, 211), (0, 0, 0), (19, 19, 19),
(40, 40, 40), (54, 54, 54), (77, 77, 77), (101, 101, 101), (129, 129, 129),
(159, 159, 159), (188, 188, 188), (226, 226, 226), (255, 255, 255)
]
self.irc_colors_tensor = torch.tensor(self.IRC_COLORS, dtype=torch.float32)
def manhattan_distance(self, detected_color, target_colors):
"""Manhattan distance function matching the color transfer node"""
"""Manhattan distance function for color matching."""
return np.sum(np.abs(detected_color - target_colors), axis=1)
def apply_color_transfer(self, image, use_16_colors=False):
"""Apply color transfer using KMeans clustering and Manhattan distance matching"""
"""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)
# 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
target_colors = self.IRC_COLORS[:16] if use_16_colors else self.IRC_COLORS
# KMeans clustering
clustering_model = KMeans(n_clusters=len(target_colors), n_init="auto", random_state=42)
@@ -80,8 +54,7 @@ class IRCArtConverter:
for color in detected_colors:
distances = self.manhattan_distance(color, target_colors_array)
closest_color = target_colors_array[np.argmin(distances)]
closest_colors.append(closest_color)
closest_colors.append(target_colors_array[np.argmin(distances)])
closest_colors = np.array(closest_colors)
@@ -91,18 +64,12 @@ class IRCArtConverter:
# Convert back to torch tensor
return torch.from_numpy(processed_image.astype(np.float32) / 255.0)
def get_irc_color_code(self, color_idx):
"""Convert IRC color array index to mIRC color code (1-99)"""
return color_idx + 1
def get_block_color(self, block, method):
"""Extract representative color from a block using specified method"""
"""Extract representative color from a block using specified method."""
if method == "average":
# Simple average of all pixels in the block
return tuple(block.mean(dim=(0, 1)).cpu().numpy() * 255)
elif method == "median":
# Median of all pixels in the block
block_flat = block.reshape(-1, 3)
median_vals = torch.median(block_flat, dim=0)[0]
return tuple(median_vals.cpu().numpy() * 255)
@@ -113,7 +80,7 @@ class IRCArtConverter:
if len(block_flat) < 2:
return tuple(block_flat[0] * 255)
# Check for color variance - if too uniform, just return average
# Check for color variance
color_std = np.std(block_flat, axis=0)
if np.max(color_std) < 0.01: # Very uniform color
return tuple(np.mean(block_flat, axis=0) * 255)
@@ -144,6 +111,76 @@ class IRCArtConverter:
# Fallback to average if k-means fails
return tuple(block.mean(dim=(0, 1)).cpu().numpy() * 255)
def find_irc_color_index(self, block_color, use_16_colors):
"""Find the IRC color index for a given block color."""
# 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
# 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)
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.manhattan_distance(block_color_array, target_colors_array)
color_idx = np.argmin(distances)
return color_idx
def format_irc_line(self, line_colors):
"""Format a line of IRC color codes with compression."""
if not line_colors:
return ""
line_text = ""
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
if color_code <= 9:
color_format = f"\x03{color_code},{color_code}"
else:
color_format = f"\x03{color_code:02d},{color_code:02d}"
line_text += color_format + (" " * count)
i += count
return line_text
def process_block(self, color_transferred_img, x, y, block_width, block_height,
actual_block_height, width, height, color_method, use_16_colors):
"""Process a single block and return its color index."""
# Calculate block boundaries
start_x = x * block_width
end_x = min(start_x + block_width, width)
start_y = round(y * actual_block_height)
end_y = round(min((y + 1) * actual_block_height, height))
# Extract block
block = color_transferred_img[start_y:end_y, start_x:end_x, :]
# Get IRC color index
if block.numel() == 0:
return 0, start_y, end_y
else:
block_color = self.get_block_color(block, color_method)
return self.find_irc_color_index(block_color, use_16_colors), start_y, end_y
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
@@ -154,102 +191,75 @@ class IRCArtConverter:
# Get the first (and only) image from batch
img = image[0] # Shape: (height, width, 3)
# Step 1: Apply color transfer to map image to IRC palette
# Apply color transfer to map image to IRC palette
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
# Calculate number of blocks
blocks_x = width // block_width
blocks_y = int(height // actual_block_height)
if half_block_mode:
blocks_y = (height * 2) // block_height
actual_block_height = block_height / 2
else:
blocks_y = height // block_height
actual_block_height = block_height
# Create output image and text lines
output = torch.zeros_like(img)
text_lines = []
# Process each block
for y in range(blocks_y):
line_colors = [] # Store color codes for this line
for x in range(blocks_x):
# Calculate block boundaries
start_x = x * block_width
end_x = min(start_x + block_width, width)
# Process blocks
if half_block_mode:
# Half-block mode: process in pairs
for y in range(0, blocks_y, 2):
line_colors = []
if half_block_mode:
start_y = round(y * actual_block_height)
end_y = round(min((y + 1) * actual_block_height, height))
else:
start_y = y * block_height
end_y = min(start_y + block_height, height)
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, color_method, use_16_colors
)
# 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 + 1) # IRC codes are 1-based
# 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, color_method, use_16_colors
)
# 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
# Extract block from color-transferred image
block = color_transferred_img[start_y:end_y, start_x:end_x, :]
text_lines.append(self.format_irc_line(line_colors))
else:
# Full block mode
for y in range(blocks_y):
line_colors = []
if block.numel() == 0:
line_colors.append(1) # Default to black
continue
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, color_method, use_16_colors
)
# 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 + 1) # IRC codes are 1-based
# Get representative color for this block (now already in IRC palette)
block_color = self.get_block_color(block, color_method)
# Find the exact IRC color index since colors are already mapped
block_color_array = np.array(block_color)
# 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)
# 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)
# Fill the entire block with the IRC color
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 optimized color codes and compression
line_text = ""
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)
text_lines.append(self.format_irc_line(line_colors))
# Join all lines with newlines (no header timestamp)
# Join all lines with newlines
irc_text = "\n".join(text_lines)
# Return as batch of 1