781 lines
27 KiB
Python
781 lines
27 KiB
Python
import gc
|
|
from itertools import combinations
|
|
import ast
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
import ot
|
|
|
|
from sklearn.cluster import KMeans, MiniBatchKMeans
|
|
from scipy.spatial import Delaunay
|
|
from skimage import color
|
|
from scipy.interpolate import Rbf
|
|
|
|
from comfy.comfy_types import IO, ComfyNodeABC
|
|
|
|
from .utils import (
|
|
EuclideanDistance,
|
|
ManhattanDistance,
|
|
CosineSimilarity,
|
|
HSVColorSimilarity,
|
|
RGBWeightedDistance,
|
|
RGBWeightedSimilarity,
|
|
Blur,
|
|
)
|
|
|
|
|
|
class PaletteExtension:
|
|
MAX_PALETTE_SIZE = 512
|
|
|
|
@staticmethod
|
|
def dense_palette(
|
|
base_palette: list[tuple[int, int, int]],
|
|
points: int = 5,
|
|
iterations: int = 2,
|
|
extend_bw: bool = True,
|
|
) -> list[tuple[int, int, int]]:
|
|
"""
|
|
Interpolate N points between each pair of colors in the base_palette.
|
|
"""
|
|
if extend_bw:
|
|
base_palette = set(base_palette)
|
|
base_palette.add((0, 0, 0))
|
|
base_palette.add((255, 255, 255))
|
|
|
|
palette = np.array(list(base_palette), dtype=np.float32)
|
|
|
|
for _ in range(iterations):
|
|
if len(palette) > PaletteExtension.MAX_PALETTE_SIZE:
|
|
break
|
|
|
|
idx_pairs = list(combinations(range(len(palette)), 2))
|
|
t = np.linspace(0, 1, points + 2)[1:-1].astype(np.float32)
|
|
|
|
new_colors = []
|
|
for i, j in idx_pairs:
|
|
c1, c2 = palette[i], palette[j]
|
|
inter = c1[None, :] + (c2 - c1)[None, :] * t[:, None]
|
|
new_colors.append(inter)
|
|
if new_colors:
|
|
new_colors = np.vstack(new_colors)
|
|
palette = np.vstack([palette, new_colors])
|
|
|
|
palette = np.unique(np.rint(palette).astype(int), axis=0).astype(np.float32)
|
|
|
|
result = [tuple(rgb.astype(int)) for rgb in palette]
|
|
return result
|
|
|
|
@staticmethod
|
|
def edge_based_palette(
|
|
base_palette: list[tuple[int, int, int]], points: int = 5, iterations: int = 2
|
|
) -> list[tuple[int, int, int]]:
|
|
"""
|
|
Use Delaunay triangulation to find edges between colors.
|
|
"""
|
|
base_palette = set(base_palette)
|
|
base_palette.add((0, 0, 0))
|
|
base_palette.add((255, 255, 255))
|
|
|
|
palette = np.array(list(base_palette), dtype=np.float32)
|
|
|
|
for _ in range(iterations):
|
|
if len(palette) > PaletteExtension.MAX_PALETTE_SIZE:
|
|
break
|
|
|
|
if len(palette) >= palette.shape[1] + 1:
|
|
try:
|
|
tri = Delaunay(palette)
|
|
edges = set()
|
|
for simplex in tri.simplices:
|
|
for i in range(len(simplex)):
|
|
for j in range(i + 1, len(simplex)):
|
|
a, b = simplex[i], simplex[j]
|
|
edges.add(tuple(sorted((a, b))))
|
|
except Exception:
|
|
# Fallback if Delaunay fails (e.g. coplanar points)
|
|
idxs = range(len(palette))
|
|
edges = set(tuple(sorted((i, j))) for i in idxs for j in idxs if i < j)
|
|
else:
|
|
idxs = range(len(palette))
|
|
edges = set(tuple(sorted((i, j))) for i in idxs for j in idxs if i < j)
|
|
|
|
new_colors = []
|
|
t = np.linspace(0, 1, points + 2)[1:-1][:, None].astype(np.float32)
|
|
|
|
for i, j in edges:
|
|
c1, c2 = palette[i], palette[j]
|
|
inters = c1 + (c2 - c1) * t
|
|
new_colors.append(inters)
|
|
|
|
if new_colors:
|
|
new_stack = np.vstack(new_colors)
|
|
palette = np.vstack([palette, new_stack])
|
|
palette = np.unique(np.rint(palette).astype(int), axis=0).astype(np.float32)
|
|
|
|
return [tuple(c.astype(int)) for c in palette]
|
|
|
|
|
|
class ColorSpaceConvert:
|
|
@staticmethod
|
|
def convert_to_target_space(
|
|
image: np.ndarray, target_colors: list[tuple[int, int, int]], color_space: str
|
|
) -> tuple[np.ndarray, list[tuple[int, int, int]]]:
|
|
if color_space == "RGB":
|
|
return image, target_colors
|
|
|
|
conversion_map = {
|
|
"HSV": (cv2.COLOR_RGB2HSV, cv2.COLOR_HSV2RGB),
|
|
"LAB": (cv2.COLOR_RGB2LAB, cv2.COLOR_LAB2RGB),
|
|
}
|
|
|
|
forward_conversion, _ = conversion_map[color_space]
|
|
|
|
# Ensure image is float32 compliant for OpenCV conversions if needed,
|
|
# though cv2 usually handles standard types. keeping input type is safer.
|
|
converted_image = cv2.cvtColor(image, forward_conversion)
|
|
|
|
target_colors_array = np.array(target_colors, dtype=np.uint8).reshape(-1, 1, 3)
|
|
converted_colors = cv2.cvtColor(target_colors_array, forward_conversion)
|
|
converted_colors = [tuple(color[0]) for color in converted_colors]
|
|
|
|
return converted_image, converted_colors
|
|
|
|
@staticmethod
|
|
def convert_to_rgb(image: np.ndarray, color_space: str) -> np.ndarray:
|
|
if color_space == "RGB":
|
|
return image
|
|
|
|
conversion_map = {"HSV": cv2.COLOR_HSV2RGB, "LAB": cv2.COLOR_LAB2RGB}
|
|
# Ensure uint8 for consistent cvtColor behavior across color spaces
|
|
return cv2.cvtColor(image.clip(0, 255).astype(np.uint8), conversion_map[color_space])
|
|
|
|
|
|
class ColorClustering:
|
|
def __init__(self, cluster_method: str):
|
|
self.clustering_methods = {
|
|
"Kmeans": KMeans,
|
|
"Mini batch Kmeans": MiniBatchKMeans,
|
|
}
|
|
self.method_name = cluster_method
|
|
self.method = self.clustering_methods[cluster_method]
|
|
|
|
def cluster_colors(self, image: np.ndarray, k: int) -> dict:
|
|
"""Perform color clustering on the image."""
|
|
# Ensure we are working with float32 to save memory
|
|
img_array = image.reshape((-1, 3)).astype(np.float32)
|
|
|
|
# Add explicit parameters to control memory usage
|
|
if self.method_name == "Mini batch Kmeans":
|
|
clustering_model = self.method(
|
|
n_clusters=k,
|
|
n_init="auto",
|
|
batch_size=4096, # Process in chunks
|
|
reassignment_ratio=0 # Reduce bookkeeping
|
|
)
|
|
else:
|
|
clustering_model = self.method(n_clusters=k, n_init="auto")
|
|
|
|
clustering_model.fit(img_array)
|
|
|
|
return {
|
|
"image": image,
|
|
"main_colors": clustering_model.cluster_centers_.astype(int),
|
|
"model": clustering_model,
|
|
}
|
|
|
|
|
|
class ColorMatcher:
|
|
def __init__(self, distance_method: str):
|
|
self.distance_methods = {
|
|
"Euclidean": EuclideanDistance,
|
|
"Manhattan": ManhattanDistance,
|
|
"Cosine Similarity": CosineSimilarity,
|
|
"HSV Distance": HSVColorSimilarity,
|
|
"RGB Weighted Distance": RGBWeightedDistance,
|
|
"RGB Weighted Similarity": RGBWeightedSimilarity,
|
|
}
|
|
self.distance_func = self.distance_methods[distance_method]
|
|
|
|
def match_colors(
|
|
self,
|
|
detected_colors: np.ndarray,
|
|
target_colors: list[tuple[int, int, int]],
|
|
clustering_model: KMeans | MiniBatchKMeans,
|
|
image_shape: tuple[int, int, int],
|
|
) -> np.ndarray:
|
|
"""Match detected colors with target colors."""
|
|
closest_colors = []
|
|
|
|
for color in detected_colors:
|
|
distances = self.distance_func(color, target_colors)
|
|
closest_color = target_colors[np.argmin(distances)]
|
|
closest_colors.append(closest_color)
|
|
|
|
closest_colors = np.array(closest_colors, dtype=np.float32)
|
|
# Reconstruct image from labels
|
|
return closest_colors[clustering_model.labels_].reshape(image_shape)
|
|
|
|
|
|
class ImagePostProcessor:
|
|
def __init__(self, gaussian_blur: int = 0):
|
|
self.gaussian_blur = gaussian_blur
|
|
|
|
def process_image(self, image: np.ndarray) -> np.ndarray:
|
|
processed = np.array(image, dtype=np.float32)
|
|
|
|
if self.gaussian_blur:
|
|
processed = Blur(processed, self.gaussian_blur)
|
|
|
|
return processed / 255.0
|
|
|
|
|
|
def process_image_with_palette(
|
|
image: list[torch.Tensor],
|
|
target_colors: list[tuple[int, int, int]],
|
|
color_space: str,
|
|
cluster_method: str,
|
|
distance_method: str,
|
|
gaussian_blur: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Shared function to process an image with the given parameters.
|
|
"""
|
|
processedImages = []
|
|
|
|
converter = ColorSpaceConvert()
|
|
clustering_engine = ColorClustering(cluster_method)
|
|
color_matcher = ColorMatcher(distance_method)
|
|
image_processor = ImagePostProcessor(gaussian_blur)
|
|
|
|
for img_tensor in image:
|
|
# Convert to uint8 0-255 so cv2.cvtColor handles all color spaces correctly
|
|
img = (255.0 * img_tensor.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
|
|
|
converted_img, converted_colors = converter.convert_to_target_space(
|
|
img, target_colors, color_space
|
|
)
|
|
|
|
clustering_result = clustering_engine.cluster_colors(
|
|
converted_img, len(target_colors)
|
|
)
|
|
|
|
processed = color_matcher.match_colors(
|
|
clustering_result["main_colors"],
|
|
converted_colors,
|
|
clustering_result["model"],
|
|
converted_img.shape,
|
|
)
|
|
|
|
processed = converter.convert_to_rgb(processed, color_space)
|
|
processed = image_processor.process_image(processed)
|
|
|
|
# Ensure contiguous memory before converting to tensor
|
|
processed = np.ascontiguousarray(processed)
|
|
processed_tensor = torch.from_numpy(processed)[None,]
|
|
processedImages.append(processed_tensor)
|
|
|
|
# CLEANUP: Explicitly delete large temporary arrays and force GC
|
|
del img
|
|
del converted_img
|
|
del clustering_result
|
|
del processed
|
|
|
|
# Only collect garbage if batch size is large to avoid perf hit on single images
|
|
if len(image) > 1:
|
|
gc.collect()
|
|
|
|
return torch.cat(processedImages, dim=0)
|
|
|
|
|
|
class ReferenceTransferReinhard(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"image_reference": (IO.IMAGE,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Reference Transfer"
|
|
|
|
def color_transfer(
|
|
self, image: torch.Tensor, image_reference: torch.Tensor
|
|
) -> tuple[torch.Tensor]:
|
|
|
|
processed_images = []
|
|
|
|
target = image_reference.cpu().numpy().astype(np.float32)
|
|
if len(target.shape) == 4:
|
|
target = np.concatenate([target[i] for i in range(target.shape[0])], axis=0)
|
|
|
|
for img_tensor in image:
|
|
source = img_tensor.cpu().numpy().astype(np.float32)
|
|
|
|
source_lab = color.rgb2lab(source)
|
|
target_lab = color.rgb2lab(target)
|
|
|
|
s_mean, s_std = source_lab.mean(axis=(0, 1)), source_lab.std(axis=(0, 1))
|
|
t_mean, t_std = target_lab.mean(axis=(0, 1)), target_lab.std(axis=(0, 1))
|
|
|
|
# Avoid division by zero for uniform channels
|
|
s_std = np.maximum(s_std, 1e-6)
|
|
|
|
result_lab = (source_lab - s_mean) / s_std * t_std + t_mean
|
|
result_rgb = np.clip(color.lab2rgb(result_lab), 0, 1)
|
|
|
|
result_tensor = torch.from_numpy(result_rgb.astype(np.float32)).unsqueeze(0)
|
|
processed_images.append(result_tensor)
|
|
|
|
del source
|
|
del source_lab
|
|
del result_lab
|
|
del result_rgb
|
|
|
|
return (torch.cat(processed_images, dim=0),)
|
|
|
|
|
|
class PaletteOptimalTransportTransfer(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
"palette_extension_method": (["Dense", "Edge", "None"], {"default": "None"}),
|
|
"palette_extension_points": ("INT", {"min": 2, "max": 20, "step": 1, "default": 5}),
|
|
"blend_mode": (["Original", "Grayscale"], {"default": "Original"}),
|
|
"blend_ratio": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_colors: list[tuple[int, int, int]],
|
|
palette_extension_method: str,
|
|
palette_extension_points: int,
|
|
blend_mode: str,
|
|
blend_ratio: float,
|
|
) -> tuple[torch.Tensor]:
|
|
|
|
if palette_extension_method == "Dense":
|
|
target_colors = PaletteExtension.dense_palette(target_colors, points=palette_extension_points)
|
|
elif palette_extension_method == "Edge":
|
|
target_colors = PaletteExtension.edge_based_palette(target_colors, points=palette_extension_points)
|
|
|
|
palette = np.array(target_colors, dtype=np.float32) / 255.0
|
|
n_palette = palette.shape[0]
|
|
palette_weights = np.ones((n_palette,)) / n_palette
|
|
|
|
processed_images = []
|
|
|
|
for img_tensor in image:
|
|
source = img_tensor.cpu().numpy().astype(np.float32)
|
|
h, w, _ = source.shape
|
|
pixels = source.reshape(-1, 3)
|
|
|
|
n_source_colors = 1000
|
|
# MEMORY FIX: Reduce batch size and disable reassignment to save memory
|
|
kmeans = MiniBatchKMeans(
|
|
n_clusters=n_source_colors,
|
|
n_init="auto",
|
|
batch_size=2048,
|
|
reassignment_ratio=0
|
|
)
|
|
kmeans.fit(pixels)
|
|
source_centroids = kmeans.cluster_centers_
|
|
pixel_labels = kmeans.labels_
|
|
|
|
source_weights = np.bincount(pixel_labels) / len(pixel_labels)
|
|
|
|
cost_matrix = ot.dist(source_centroids, palette, metric="euclidean") ** 2
|
|
|
|
transport_plan = ot.sinkhorn(
|
|
source_weights,
|
|
palette_weights,
|
|
cost_matrix,
|
|
reg=1e-2,
|
|
numItermax=100000,
|
|
)
|
|
|
|
mapped_centroids = np.dot(transport_plan, palette) / np.sum(
|
|
transport_plan, axis=1, keepdims=True
|
|
)
|
|
|
|
if blend_mode == "Original":
|
|
recolored_pixels = (1 - blend_ratio) * pixels + blend_ratio * mapped_centroids[pixel_labels]
|
|
elif blend_mode == "Grayscale":
|
|
gray = color.rgb2gray(source)
|
|
gray_rgb = np.stack([gray] * 3, axis=-1)
|
|
gray_pixels = gray_rgb.reshape(-1, 3)
|
|
recolored_pixels = (1 - blend_ratio) * gray_pixels + blend_ratio * mapped_centroids[pixel_labels]
|
|
|
|
recolored_image = recolored_pixels.reshape(h, w, 3).astype(np.float32)
|
|
|
|
result_tensor = torch.from_numpy(recolored_image).unsqueeze(0)
|
|
processed_images.append(result_tensor)
|
|
|
|
# CLEANUP
|
|
del kmeans
|
|
del source
|
|
del pixels
|
|
del cost_matrix
|
|
del transport_plan
|
|
|
|
if len(image) > 1:
|
|
gc.collect()
|
|
|
|
return (torch.cat(processed_images, dim=0),)
|
|
|
|
|
|
class PaletteRbfTransfer(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
"rbf_function": (["thin_plate", "multiquadric", "inverse", "gaussian"], {"default": "gaussian"}),
|
|
"epsilon": ("FLOAT", {"min": 0.01, "max": 100, "step": 0.1, "default": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_colors: list[tuple[int, int, int]],
|
|
rbf_function: str,
|
|
epsilon: float,
|
|
) -> tuple[torch.Tensor]:
|
|
|
|
palette = np.array(target_colors, dtype=np.float32) / 255
|
|
r, g, b = palette[:, 0], palette[:, 1], palette[:, 2]
|
|
|
|
rbf_r = Rbf(r, g, b, r, function=rbf_function, epsilon=epsilon)
|
|
rbf_g = Rbf(r, g, b, g, function=rbf_function, epsilon=epsilon)
|
|
rbf_b = Rbf(r, g, b, b, function=rbf_function, epsilon=epsilon)
|
|
|
|
processed_images = []
|
|
|
|
for img_tensor in image:
|
|
source = img_tensor.cpu().numpy().astype(np.float32)
|
|
h, w, _ = source.shape
|
|
pixels = source.reshape(-1, 3)
|
|
|
|
mapped_r = rbf_r(pixels[:, 0], pixels[:, 1], pixels[:, 2])
|
|
mapped_g = rbf_g(pixels[:, 0], pixels[:, 1], pixels[:, 2])
|
|
mapped_b = rbf_b(pixels[:, 0], pixels[:, 1], pixels[:, 2])
|
|
|
|
mapped_pixels = np.stack((mapped_r, mapped_g, mapped_b), axis=-1)
|
|
mapped_pixels = np.clip(mapped_pixels, 0, 1).astype(np.float32)
|
|
recolored_image = mapped_pixels.reshape(h, w, 3)
|
|
|
|
result_tensor = torch.from_numpy(recolored_image).unsqueeze(0)
|
|
processed_images.append(result_tensor)
|
|
|
|
del source
|
|
del pixels
|
|
del mapped_pixels
|
|
|
|
return (torch.cat(processed_images, dim=0),)
|
|
|
|
|
|
class PaletteSoftTransfer(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
"blend_mode": (["Original", "Grayscale"], {"default": "Original"}),
|
|
"blend_ratio": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}),
|
|
"softness": ("FLOAT", {"min": 0, "max": 20, "step": 0.1, "default": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_colors: list[tuple[int, int, int]],
|
|
blend_mode: str,
|
|
blend_ratio: float,
|
|
softness: float,
|
|
) -> tuple[torch.Tensor]:
|
|
|
|
if len(target_colors) < 2:
|
|
raise ValueError("Palette must contain at least 2 colors")
|
|
|
|
processed_images = []
|
|
|
|
for img_tensor in image:
|
|
source = img_tensor.cpu().numpy().astype(np.float32)
|
|
|
|
img_lab = color.rgb2lab(source)
|
|
palette_lab = np.array(
|
|
[color.rgb2lab(np.array([[c]]) / 255.0)[0, 0] for c in target_colors]
|
|
)
|
|
|
|
pixels = img_lab.reshape(-1, 3).astype(np.float32)
|
|
n_pixels = pixels.shape[0]
|
|
n_colors = len(palette_lab)
|
|
chunk_size = max(1, 500_000 // n_colors)
|
|
|
|
projected = np.empty_like(pixels)
|
|
for start in range(0, n_pixels, chunk_size):
|
|
end = min(start + chunk_size, n_pixels)
|
|
chunk = pixels[start:end]
|
|
# (n_colors, chunk_len)
|
|
dists = np.array([np.linalg.norm(chunk - p, axis=1) for p in palette_lab])
|
|
weights = np.exp(-softness * dists)
|
|
weights /= weights.sum(axis=0)
|
|
projected[start:end] = np.dot(weights.T, palette_lab)
|
|
|
|
if blend_mode == "Original":
|
|
blended = (1 - blend_ratio) * pixels + blend_ratio * projected
|
|
elif blend_mode == "Grayscale":
|
|
gray = color.rgb2gray(source)
|
|
gray_rgb = np.stack([gray] * 3, axis=-1)
|
|
gray_lab = color.rgb2lab(gray_rgb)
|
|
gray_pixels = gray_lab.reshape(-1, 3)
|
|
blended = (1 - blend_ratio) * gray_pixels + blend_ratio * projected
|
|
|
|
blended_lab = blended.reshape(img_lab.shape)
|
|
blended_rgb = np.clip(color.lab2rgb(blended_lab), 0, 1).astype(np.float32)
|
|
|
|
result_tensor = torch.from_numpy(blended_rgb).unsqueeze(0)
|
|
processed_images.append(result_tensor)
|
|
|
|
del source
|
|
del img_lab
|
|
del pixels
|
|
|
|
return (torch.cat(processed_images, dim=0),)
|
|
|
|
|
|
class PaletteTransferReinhard(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self, image: torch.Tensor, target_colors: list[tuple[int, int, int]]
|
|
) -> tuple[torch.Tensor]:
|
|
if len(target_colors) == 0:
|
|
return (image,)
|
|
|
|
target_colors = PaletteExtension.dense_palette(target_colors, points=3)
|
|
|
|
# Compute target LAB stats directly from palette instead of creating a large image
|
|
palette_rgb = np.array(target_colors, dtype=np.float32).reshape(-1, 1, 3) / 255.0
|
|
palette_lab = color.rgb2lab(palette_rgb).reshape(-1, 3)
|
|
t_mean = palette_lab.mean(axis=0)
|
|
t_std = np.maximum(palette_lab.std(axis=0), 1e-6)
|
|
|
|
processed_images = []
|
|
|
|
for img_tensor in image:
|
|
source = img_tensor.cpu().numpy().astype(np.float32)
|
|
|
|
source_lab = color.rgb2lab(source)
|
|
|
|
s_mean, s_std = source_lab.mean(axis=(0, 1)), source_lab.std(axis=(0, 1))
|
|
|
|
# Avoid division by zero for uniform channels
|
|
s_std = np.maximum(s_std, 1e-6)
|
|
|
|
result_lab = (source_lab - s_mean) / s_std * t_std + t_mean
|
|
result_rgb = np.clip(color.lab2rgb(result_lab), 0, 1).astype(np.float32)
|
|
|
|
result_tensor = torch.from_numpy(result_rgb).unsqueeze(0)
|
|
processed_images.append(result_tensor)
|
|
|
|
del source
|
|
del source_lab
|
|
|
|
return (torch.cat(processed_images, dim=0),)
|
|
|
|
|
|
class PalleteTransferClustering(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
"palette_extension_method": (["Dense", "Edge", "None"], {"default": "None"}),
|
|
"palette_extension_points": ("INT", {"min": 2, "max": 20, "step": 1, "default": 5}),
|
|
"gaussian_blur": ("INT", {"default": 3, "min": 0, "max": 27, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_colors: list[tuple[int, int, int]],
|
|
palette_extension_method: str,
|
|
palette_extension_points: int,
|
|
gaussian_blur: int,
|
|
) -> tuple[torch.Tensor]:
|
|
if len(target_colors) == 0:
|
|
return (image,)
|
|
|
|
if palette_extension_method == "Dense":
|
|
target_colors = PaletteExtension.dense_palette(target_colors, points=palette_extension_points)
|
|
elif palette_extension_method == "Edge":
|
|
target_colors = PaletteExtension.edge_based_palette(target_colors, points=palette_extension_points)
|
|
|
|
output = process_image_with_palette(
|
|
image=image,
|
|
target_colors=target_colors,
|
|
color_space="RGB",
|
|
cluster_method="Mini batch Kmeans",
|
|
distance_method="Euclidean",
|
|
gaussian_blur=gaussian_blur,
|
|
)
|
|
|
|
return (output,)
|
|
|
|
|
|
class PaletteTransferNode(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"target_colors": ("COLOR_LIST",),
|
|
"color_space": (["RGB", "HSV", "LAB"], {"default": "RGB"}),
|
|
"cluster_method": (["Kmeans", "Mini batch Kmeans"], {"default": "Kmeans"}),
|
|
"distance_method": ([
|
|
"Euclidean",
|
|
"Manhattan",
|
|
"Cosine Similarity",
|
|
"HSV Distance",
|
|
"RGB Weighted Distance",
|
|
"RGB Weighted Similarity",
|
|
], {"default": "Euclidean"}),
|
|
"gaussian_blur": ("INT", {"default": 3, "min": 0, "max": 27, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
FUNCTION = "color_transfer"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_transfer(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_colors: list[tuple[int, int, int]],
|
|
color_space: str,
|
|
cluster_method: str,
|
|
distance_method: str,
|
|
gaussian_blur: int,
|
|
) -> tuple[torch.Tensor]:
|
|
if len(target_colors) == 0:
|
|
return (image,)
|
|
|
|
output = process_image_with_palette(
|
|
image=image,
|
|
target_colors=target_colors,
|
|
color_space=color_space,
|
|
cluster_method=cluster_method,
|
|
distance_method=distance_method,
|
|
gaussian_blur=gaussian_blur,
|
|
)
|
|
|
|
return (output,)
|
|
|
|
|
|
class ColorPaletteNode(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(s) -> dict:
|
|
return {
|
|
"required": {
|
|
"color_palette": (
|
|
IO.STRING,
|
|
{
|
|
"default": "[(30, 32, 30), (60, 61, 55), (105, 117, 101), (236, 223, 204)]",
|
|
"multiline": True,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("COLOR_LIST",)
|
|
RETURN_NAMES = ("Color palette",)
|
|
FUNCTION = "color_list"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def color_list(self, color_palette: str) -> tuple[list[tuple[int, int, int]]]:
|
|
parsed = ast.literal_eval(color_palette)
|
|
if not isinstance(parsed, list) or not all(
|
|
isinstance(c, tuple) and len(c) == 3 and all(isinstance(v, int) and 0 <= v <= 255 for v in c)
|
|
for c in parsed
|
|
):
|
|
raise ValueError("Color palette must be a list of (R, G, B) tuples with values 0-255")
|
|
return (parsed,)
|
|
|
|
|
|
class ExtractPaletteNode(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"image": (IO.IMAGE,),
|
|
"num_colors": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1}),
|
|
"cluster_method": (["Kmeans", "Mini batch Kmeans"], {"default": "Mini batch Kmeans"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("COLOR_LIST",)
|
|
RETURN_NAMES = ("Color palette",)
|
|
FUNCTION = "extract_palette"
|
|
CATEGORY = "Color Transfer/Palette Transfer"
|
|
|
|
def extract_palette(
|
|
self, image: torch.Tensor, num_colors: int, cluster_method: str
|
|
) -> tuple[list[tuple[int, int, int]]]:
|
|
|
|
img_tensor = image[0] if len(image.shape) == 4 else image
|
|
|
|
# MEMORY FIX: Ensure float32 for clustering
|
|
img = (255.0 * img_tensor.cpu().numpy()).astype(np.float32)
|
|
|
|
clustering_engine = ColorClustering(cluster_method)
|
|
clustering_result = clustering_engine.cluster_colors(img, num_colors)
|
|
|
|
colors = clustering_result["main_colors"]
|
|
color_list = [tuple(map(int, color)) for color in colors]
|
|
|
|
# Cleanup
|
|
del img
|
|
del clustering_result
|
|
|
|
return (color_list,) |