231 lines
8.8 KiB
Python
231 lines
8.8 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import os
|
|
|
|
class PixelPaletteExtractor:
|
|
"""
|
|
Un node ComfyUI qui extrait la palette de couleurs d'une image GIF indexée
|
|
et génère une image de palette avec option d'affichage des index.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"palette_width": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 32,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"color_size": ("INT", {
|
|
"default": 32,
|
|
"min": 8,
|
|
"max": 128,
|
|
"step": 4,
|
|
"display": "number"
|
|
}),
|
|
"show_indices": ("BOOLEAN", {"default": False}),
|
|
"font_size": ("INT", {
|
|
"default": 10,
|
|
"min": 6,
|
|
"max": 24,
|
|
"step": 1,
|
|
"display": "number"
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("palette_image",)
|
|
FUNCTION = "extract_palette"
|
|
CATEGORY = "image/color"
|
|
|
|
def extract_palette(self, image, palette_width, color_size, show_indices, font_size):
|
|
"""
|
|
Extrait la palette de couleurs de l'image et génère une image de palette
|
|
|
|
Args:
|
|
image: Tensor d'image ComfyUI (batch, height, width, channels)
|
|
palette_width: Nombre de couleurs par ligne dans la palette
|
|
color_size: Taille de chaque carré de couleur en pixels
|
|
show_indices: Afficher les index des couleurs
|
|
font_size: Taille de la police pour les index
|
|
"""
|
|
|
|
try:
|
|
# Convertir le tensor ComfyUI en PIL Image
|
|
# ComfyUI utilise le format [batch, height, width, channels] avec valeurs 0-1
|
|
if len(image.shape) == 4:
|
|
img_tensor = image[0] # Prendre la première image du batch
|
|
else:
|
|
img_tensor = image
|
|
|
|
# Convertir de [0,1] à [0,255]
|
|
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
# Créer l'image PIL en RGB
|
|
if len(img_np.shape) == 3 and img_np.shape[2] == 3:
|
|
img_pil = Image.fromarray(img_np, 'RGB')
|
|
elif len(img_np.shape) == 3 and img_np.shape[2] == 4:
|
|
img_pil = Image.fromarray(img_np, 'RGBA').convert('RGB')
|
|
else:
|
|
# Fallback pour les images en niveaux de gris
|
|
if len(img_np.shape) == 2:
|
|
img_pil = Image.fromarray(img_np, 'L').convert('RGB')
|
|
else:
|
|
img_pil = Image.fromarray(img_np[:,:,0], 'L').convert('RGB')
|
|
|
|
# Extraire les couleurs uniques
|
|
colors = self.extract_unique_colors(img_pil)
|
|
|
|
# Créer l'image de palette
|
|
palette_img = self.create_palette_image(colors, palette_width, color_size, show_indices, font_size)
|
|
|
|
# Convertir l'image PIL en tensor ComfyUI
|
|
palette_array = np.array(palette_img).astype(np.float32) / 255.0
|
|
|
|
# S'assurer que le tensor a la bonne forme [batch, height, width, channels]
|
|
if len(palette_array.shape) == 3:
|
|
palette_tensor = torch.from_numpy(palette_array).unsqueeze(0)
|
|
else:
|
|
palette_tensor = torch.from_numpy(palette_array)
|
|
|
|
return (palette_tensor,)
|
|
|
|
except Exception as e:
|
|
print(f"Erreur dans PixelPaletteExtractor: {e}")
|
|
# Retourner une image d'erreur
|
|
error_img = Image.new('RGB', (color_size, color_size), (255, 0, 0))
|
|
error_array = np.array(error_img).astype(np.float32) / 255.0
|
|
error_tensor = torch.from_numpy(error_array).unsqueeze(0)
|
|
return (error_tensor,)
|
|
|
|
def extract_unique_colors(self, img_pil):
|
|
"""Extrait les couleurs uniques de l'image"""
|
|
colors = []
|
|
|
|
try:
|
|
# Méthode 1: Essayer de convertir en palette indexée
|
|
img_palette = img_pil.convert('P', colors=256)
|
|
palette_data = img_palette.getpalette()
|
|
|
|
if palette_data is not None:
|
|
# Extraire les couleurs uniques utilisées
|
|
img_array = np.array(img_palette)
|
|
unique_indices = np.unique(img_array)
|
|
|
|
for idx in unique_indices:
|
|
if idx * 3 + 2 < len(palette_data):
|
|
r = palette_data[idx * 3]
|
|
g = palette_data[idx * 3 + 1]
|
|
b = palette_data[idx * 3 + 2]
|
|
colors.append((r, g, b, idx))
|
|
else:
|
|
raise Exception("Pas de palette trouvée")
|
|
|
|
except:
|
|
# Méthode 2: Quantifier l'image directement
|
|
try:
|
|
img_quantized = img_pil.quantize(colors=256)
|
|
palette_data = img_quantized.getpalette()
|
|
|
|
if palette_data is not None:
|
|
img_array = np.array(img_quantized)
|
|
unique_indices = np.unique(img_array)
|
|
|
|
for idx in unique_indices:
|
|
if idx * 3 + 2 < len(palette_data):
|
|
r = palette_data[idx * 3]
|
|
g = palette_data[idx * 3 + 1]
|
|
b = palette_data[idx * 3 + 2]
|
|
colors.append((r, g, b, idx))
|
|
else:
|
|
raise Exception("Quantification échouée")
|
|
|
|
except:
|
|
# Méthode 3: Extraire les couleurs directement (fallback)
|
|
img_array = np.array(img_pil)
|
|
# Reshaper et obtenir les couleurs uniques
|
|
pixels = img_array.reshape(-1, 3)
|
|
unique_colors = np.unique(pixels, axis=0)
|
|
|
|
for i, color in enumerate(unique_colors):
|
|
colors.append((int(color[0]), int(color[1]), int(color[2]), i))
|
|
|
|
return colors
|
|
|
|
def create_palette_image(self, colors, palette_width, color_size, show_indices, font_size):
|
|
"""Crée l'image de palette"""
|
|
num_colors = len(colors)
|
|
|
|
if num_colors == 0:
|
|
# Image vide, retourner une image noire
|
|
return Image.new('RGB', (color_size, color_size), (0, 0, 0))
|
|
|
|
palette_height = (num_colors + palette_width - 1) // palette_width
|
|
|
|
# Taille de l'image finale
|
|
img_width = palette_width * color_size
|
|
img_height = palette_height * color_size
|
|
|
|
# Créer l'image de palette
|
|
palette_img = Image.new('RGB', (img_width, img_height), (255, 255, 255))
|
|
draw = ImageDraw.Draw(palette_img)
|
|
|
|
# Charger une police si possible
|
|
font = None
|
|
if show_indices:
|
|
try:
|
|
font = ImageFont.load_default()
|
|
except:
|
|
font = None
|
|
|
|
# Dessiner chaque couleur
|
|
for i, (r, g, b, idx) in enumerate(colors):
|
|
row = i // palette_width
|
|
col = i % palette_width
|
|
|
|
x1 = col * color_size
|
|
y1 = row * color_size
|
|
x2 = x1 + color_size
|
|
y2 = y1 + color_size
|
|
|
|
# Dessiner le carré de couleur
|
|
draw.rectangle([x1, y1, x2-1, y2-1], fill=(r, g, b))
|
|
|
|
# Ajouter l'index si demandé
|
|
if show_indices and font is not None:
|
|
text = str(idx)
|
|
|
|
# Position du texte (centré approximativement)
|
|
text_x = x1 + 4
|
|
text_y = y1 + 4
|
|
|
|
# Fond blanc pour le texte
|
|
try:
|
|
bbox = draw.textbbox((text_x, text_y), text, font=font)
|
|
draw.rectangle([bbox[0]-1, bbox[1]-1, bbox[2]+1, bbox[3]+1],
|
|
fill=(255, 255, 255))
|
|
except:
|
|
# Fallback si textbbox n'est pas disponible
|
|
draw.rectangle([text_x-1, text_y-1, text_x+20, text_y+12],
|
|
fill=(255, 255, 255))
|
|
|
|
# Dessiner le texte en noir
|
|
draw.text((text_x, text_y), text, fill=(0, 0, 0), font=font)
|
|
|
|
return palette_img
|
|
|
|
# Mappage des nodes pour ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PixelPaletteExtractor": PixelPaletteExtractor
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PixelPaletteExtractor": "Pixel Palette Extractor"
|
|
}
|