nodes/852wa-ComfyUI-ColorshiftColor
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

塗りつぶし影分離ノードを追加

Browse Source
This commit is contained in:
hakoniwa
2025-02-01 21:07:49 +09:00
committed by GitHub
parent d7a34b647d
commit a280bf9c86
1 changed files with 142 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+142 -131
View File
@@ -1,7 +1,5 @@
# nodes.py
import json
import random
import numpy as np
import torch
import torch.nn.functional as F
@@ -9,7 +7,6 @@ from PIL import Image, ImageDraw, ImageFont
from scipy.cluster.vq import vq
from sklearn.cluster import KMeans
class ColorshiftColorNode:
@classmethod
def INPUT_TYPES(cls):
@@ -32,72 +29,50 @@ class ColorshiftColorNode:
FUNCTION = "process"
CATEGORY = "Image/Color"
def process(
self,
images,
color_count,
lock_masks=None,
palette_override=None,
font_size=20,
sampling_rate=0.25,
n_init=3,
):
# デバイスを取得
device = images.device
batch_size, height, width, channels = images.shape # imagesの形状を使用
def process(self, images, color_count, lock_masks=None, palette_override=None, font_size=20, sampling_rate=0.25, n_init=3):
batch_size, height, width, channels = images.shape
pixels = images.reshape(batch_size, -1, 3)
if palette_override is not None:
palette = palette_override.to(device)
# バッチ全体に対して一回の距離計算でラベルを取得
palette = palette_override
labels = self._match_palette_batch(pixels, palette)
else:
# 複数の画像に対して、それぞれでクラスタリングをすると、計算量も増えるし、パレットがずれる。
# 解決策として、複数の画像を混ぜてクラスタリングを行う
# 混ぜてクラスタリング
# パレット生成時にのみリサイズ
resized_images = F.interpolate(
images.permute(0, 3, 1, 2),
scale_factor=0.5,
mode="bilinear",
align_corners=False,
).permute(0, 2, 3, 1)
pixels_for_clustering = (
resized_images.reshape(batch_size, -1, 3).cpu().numpy().reshape(-1, 3)
)
random_indices = np.random.choice(
pixels_for_clustering.shape[0],
size=int(pixels_for_clustering.shape[0] * sampling_rate),
replace=False,
)
# パレット生成のために画像をリサイズ
resized_images = F.interpolate(images.permute(0, 3, 1, 2), scale_factor=0.5, mode="bilinear", align_corners=False).permute(0, 2, 3, 1)
pixels_for_clustering = resized_images.reshape(batch_size, -1, 3).cpu().numpy().reshape(-1, 3)
random_indices = np.random.choice(pixels_for_clustering.shape[0], size=int(pixels_for_clustering.shape[0] * sampling_rate), replace=False)
pixels_for_clustering = pixels_for_clustering[random_indices]
kmeans = KMeans(n_clusters=color_count, random_state=42, n_init=n_init)
kmeans.fit(pixels_for_clustering)
# バッチ内の各画像に対してラベルを取得 (scipy.cluster.vq.vq を使用)
labels, dist = vq(
pixels.cpu().numpy().reshape(-1, 3), kmeans.cluster_centers_
)
labels = torch.from_numpy(labels.reshape(batch_size, -1)).long().to(device)
palette = torch.from_numpy(kmeans.cluster_centers_).float().to(device)
# 全画素に対してクラスタ中心からの最近傍を求める
labels, dist = vq(pixels.cpu().numpy().reshape(-1, 3), kmeans.cluster_centers_)
labels = torch.from_numpy(labels.reshape(batch_size, -1)).long()
palette = torch.from_numpy(kmeans.cluster_centers_).float()
# ----- パレットの並び順を各クラスタの占有率が高い順に変更 -----
# 全ラベルを1次元にまとめ、各クラスタの出現回数を計算
all_labels = labels.flatten()
counts = torch.bincount(all_labels, minlength=color_count)
# 出現回数の降順に並べ替えたときの各クラスタの元のインデックス
sorted_order = torch.argsort(counts, descending=True)
# mapping: 旧インデックス -> 新インデックス を作成
mapping = torch.empty_like(sorted_order)
mapping[sorted_order] = torch.arange(sorted_order.size(0), device=sorted_order.device)
# パレットを並び替え、各画素のラベルを新しい順序に置換
palette = palette[sorted_order]
labels = mapping[labels]
# ---------------------------------------------------------------
# Generate palette preview (最初の画像のみ, CPUで処理)
preview_img = self.generate_palette_preview(palette.cpu().numpy(), font_size)
# バッチ全体のインデックスマップを取得
index_maps = labels.reshape(batch_size, height, width)
# GPUで一括処理
if lock_masks is not None:
lock_masks_batch = []
for i in range(batch_size):
lock_mask = lock_masks[i] if i < lock_masks.shape[0] else lock_masks[-1]
lock_masks_batch.append(lock_mask)
lock_masks = torch.stack(lock_masks_batch).to(device)
lock_masks = torch.stack(lock_masks_batch)
pixel_masks = torch.zeros_like(index_maps, dtype=torch.bool)
for i in range(batch_size):
@@ -105,7 +80,7 @@ class ColorshiftColorNode:
if is_locked:
pixel_masks[i][index_maps[i] == j] = True
original_pixels = pixels.clone() # ディープコピー
original_pixels = pixels.clone()
new_pixels = palette[labels]
pixels[~pixel_masks] = new_pixels[~pixel_masks]
pixels[pixel_masks] = original_pixels[pixel_masks]
@@ -137,32 +112,20 @@ class ColorshiftColorNode:
x = (i % cols) * patch_size
y = (i // cols) * patch_size
draw.rectangle(
[x, y, x + patch_size, y + patch_size],
fill=tuple((color * 255).astype(int)),
)
draw.rectangle([x, y, x + patch_size, y + patch_size], fill=tuple((color * 255).astype(int)))
text = str(i)
bbox = draw.textbbox((x, y), text, font=font)
text_w = bbox[2] - bbox[0]
text_h = bbox[3] - bbox[1]
draw.text(
(x + (patch_size - text_w) // 2, y + (patch_size - text_h) // 2),
text,
fill="white",
font=font,
)
draw.text((x + (patch_size - text_w) // 2, y + (patch_size - text_h) // 2), text, fill="white", font=font)
img_tensor = torch.from_numpy(
np.array(img).astype(np.float32) / 255.0
).unsqueeze(0)
img_tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
return img_tensor
def _match_palette_batch(self, pixels, palette):
distances = torch.cdist(pixels, palette)
return torch.argmin(distances, dim=2)
class PaletteEditorNode:
@classmethod
def INPUT_TYPES(cls):
@@ -180,10 +143,7 @@ class PaletteEditorNode:
"saturation_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0}),
"value_random_enable": ("BOOLEAN", {"default": False}),
"value_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0}),
"lock_color_num": (
"STRING",
{"default": "0,", "placeholder": "例: 0,1,2,3,4"},
),
"lock_color_num": ("STRING", {"default": "0,", "placeholder": "例: 0,1,2,3,4"}),
"mask_enable": ("BOOLEAN", {"default": False}),
"invert_mask": ("BOOLEAN", {"default": False}),
"operations": ("STRING", {"default": "[]", "multiline": True}),
@@ -250,76 +210,43 @@ class PaletteEditorNode:
rgb = rgb_prime + m[..., None]
return rgb
def process_palette(
self,
images,
palette,
index_maps,
seed,
hue_random_enable,
hue_shift,
saturation_random_enable,
saturation_scale,
value_random_enable,
value_scale,
lock_color_num="0,",
mask_enable=False,
invert_mask=False,
operations="[]",
):
device = images.device
palette = palette.to(device)
index_maps = index_maps.to(device)
def process_palette(self, images, palette, index_maps, seed, hue_random_enable, hue_shift, saturation_random_enable, saturation_scale, value_random_enable, value_scale, lock_color_num="0,", mask_enable=False, invert_mask=False, operations="[]"):
palette = palette
index_maps = index_maps
batch_size = images.shape[0]
height = images.shape[1]
width = images.shape[2]
if not mask_enable:
color_mask = torch.zeros(len(palette), dtype=torch.bool, device=device)
color_mask = torch.zeros(len(palette), dtype=torch.bool)
else:
try:
indices = [
int(idx.strip()) for idx in lock_color_num.split(",") if idx.strip()
]
color_mask = torch.zeros(len(palette), dtype=torch.bool, device=device)
indices = [int(idx.strip()) for idx in lock_color_num.split(",") if idx.strip()]
color_mask = torch.zeros(len(palette), dtype=torch.bool)
for idx in indices:
if 0 <= idx < len(palette):
color_mask[idx] = True
if invert_mask:
color_mask = ~color_mask
except ValueError:
color_mask = torch.zeros(len(palette), dtype=torch.bool, device=device)
color_mask = torch.zeros(len(palette), dtype=torch.bool)
current_hue = random.uniform(-180, 180) if hue_random_enable else hue_shift
current_saturation = (
random.uniform(0, 5) if saturation_random_enable else saturation_scale
)
current_saturation = random.uniform(0, 5) if saturation_random_enable else saturation_scale
current_value = random.uniform(0, 5) if value_random_enable else value_scale
modified_palette = self.edit_palette(
palette,
operations,
color_mask,
current_hue,
current_saturation,
current_value,
)
modified_palette = self.edit_palette(palette, operations, color_mask, current_hue, current_saturation, current_value)
# バッチ処理
pixels = images.reshape(batch_size, -1, 3)
labels = index_maps.reshape(batch_size, -1)
new_pixels = modified_palette[labels.long()]
original_pixels = pixels.clone()
# ピクセルマスクの生成（2D形式）
pixel_mask = torch.zeros(
(batch_size, height, width), dtype=torch.float32, device=device
)
pixel_mask = torch.zeros((batch_size, height, width), dtype=torch.float32)
for i, is_masked in enumerate(color_mask):
pixel_mask[index_maps == i] = 1.0 if is_masked else 0.0
# 画像の処理
pixels_mask = pixel_mask.reshape(batch_size, -1, 1).expand(-1, -1, 3)
processed_pixels = torch.where(pixels_mask > 0.5, original_pixels, new_pixels)
processed_images = processed_pixels.reshape(images.shape)
@@ -327,15 +254,11 @@ class PaletteEditorNode:
random.seed(None)
return (processed_images, modified_palette, pixel_mask)
def edit_palette(
self, palette, operations, mask, hue_shift, saturation_scale, value_scale
):
def edit_palette(self, palette, operations, mask, hue_shift, saturation_scale, value_scale):
hsv_palette = self.rgb_to_hsv(palette)
hsv_palette[~mask, 0] = (hsv_palette[~mask, 0] + hue_shift / 360.0) % 1.0
hsv_palette[~mask, 1] = torch.clamp(
hsv_palette[~mask, 1] * saturation_scale, 0, 1
)
hsv_palette[~mask, 1] = torch.clamp(hsv_palette[~mask, 1] * saturation_scale, 0, 1)
hsv_palette[~mask, 2] = torch.clamp(hsv_palette[~mask, 2] * value_scale, 0, 1)
try:
@@ -346,32 +269,120 @@ class PaletteEditorNode:
idx = op["index"]
if 0 <= idx < len(palette):
if "color" in op and len(op["color"]) == 3:
new_color = torch.tensor(
op["color"],
dtype=torch.float32,
device=palette.device,
)
new_color = torch.tensor(op["color"], dtype=torch.float32)
palette[idx] = new_color
hsv_palette[idx] = self.rgb_to_hsv(
new_color.reshape(1, 3)
)[0]
hsv_palette[idx] = self.rgb_to_hsv(new_color.reshape(1, 3))[0]
if "hsv" in op and len(op["hsv"]) == 3:
hsv_palette[idx] = torch.clamp(
torch.tensor(op["hsv"], device=palette.device), 0, 1
)
hsv_palette[idx] = torch.clamp(torch.tensor(op["hsv"]), 0, 1)
except Exception as e:
print(f"Operation Error: {str(e)}")
modified_palette = self.hsv_to_rgb(hsv_palette)
return modified_palette
class CsCFill:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"palette": ("PALETTE",),
"index_maps": ("MASK",),
},
"optional": {
"operations": ("STRING", {"default": "[]", "multiline": True}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE")
RETURN_NAMES = ("filled_image", "mask_image")
FUNCTION = "process"
CATEGORY = "Image/Color"
def process(self, images, palette, index_maps, operations="[]"):
batch_size, height, width, channels = images.shape
pairs = None
try:
ops = json.loads(operations)
if isinstance(ops, list) and len(ops) > 0:
manual_pairs = []
for op in ops:
if isinstance(op, dict) and "A" in op and "B" in op:
try:
a_idx = int(op["A"])
b_idx = int(op["B"])
if 0 <= a_idx < palette.shape[0] and 0 <= b_idx < palette.shape[0]:
brightness_a = palette[a_idx].sum()
brightness_b = palette[b_idx].sum()
if brightness_a < brightness_b:
a_idx, b_idx = b_idx, a_idx
manual_pairs.append((a_idx, b_idx))
except Exception:
continue
if len(manual_pairs) > 0:
pairs = manual_pairs
except Exception as e:
pairs = None
if pairs is None:
pairs = self._compute_auto_pairs(palette)
filled_image = images.clone()
mask_image = torch.zeros(batch_size, height, width, 4, dtype=images.dtype)
for (a_idx, b_idx) in pairs:
mask = (index_maps == b_idx)
mask_expanded_rgb = mask.unsqueeze(-1).expand(-1, -1, -1, 3)
filled_color = palette[a_idx].view(1, 1, 1, 3).expand_as(filled_image)
filled_image = torch.where(mask_expanded_rgb, filled_color, filled_image)
b_color = palette[b_idx].view(1, 1, 1, 3).expand(batch_size, height, width, 3)
current_rgb = mask_image[..., :3]
updated_rgb = torch.where(mask_expanded_rgb, b_color, current_rgb)
mask_image[..., :3] = updated_rgb
current_alpha = mask_image[..., 3:4]
new_alpha = torch.where(mask.unsqueeze(-1), torch.tensor(1.0, dtype=images.dtype), current_alpha)
mask_image[..., 3:4] = new_alpha
return (filled_image, mask_image)
def _compute_auto_pairs(self, palette):
palette_np = palette.detach().cpu().numpy()
n = palette_np.shape[0]
pair_list = []
for i in range(n):
for j in range(i + 1, n):
dist = np.linalg.norm(palette_np[i] - palette_np[j])
pair_list.append((i, j, dist))
pair_list.sort(key=lambda x: x[2])
paired = set()
auto_pairs = []
for i, j, dist in pair_list:
if i in paired or j in paired:
continue
brightness_i = palette_np[i].sum()
brightness_j = palette_np[j].sum()
if brightness_i >= brightness_j:
auto_pairs.append((i, j))
else:
auto_pairs.append((j, i))
paired.add(i)
paired.add(j)
return auto_pairs
NODE_CLASS_MAPPINGS = {
"ColorshiftColor": ColorshiftColorNode,
"CsCPaletteEditor": PaletteEditorNode,
"CsCFill": CsCFill,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ColorshiftColor": "ColorshiftColor",
"CsCPaletteEditor": "CsCPaletteEditor",
"CsCFill": "CsCFill",
}