| import modules.scripts as scripts
|
| import gradio as gr
|
|
|
| from modules import images
|
| from modules.processing import process_images
|
| from modules.shared import opts
|
| import numpy as np
|
|
|
|
|
| class Script(scripts.Script):
|
|
|
| def title(self):
|
| return "txt2palette"
|
|
|
| def show(self, is_img2img):
|
| return not is_img2img
|
|
|
| def ui(self, is_img2img):
|
| palette_size = gr.Slider(minimum=1, maximum=64, step=1, value=0,
|
| label="Palette size")
|
| method = gr.Radio(choices=['Median cut', 'KMeans'], value='Median cut', label='Palette extraction method')
|
| sort_by = gr.Radio(choices=["luminance", "hue", "saturation", "value", "lightness"], value="luminance", label="Sort colors by")
|
| overwrite = gr.Checkbox(False, label="Overwrite existing files")
|
| return [palette_size, method, sort_by, overwrite]
|
|
|
| def run(self, p, palette_size, method, sort_by, overwrite):
|
| import colorsys
|
| from PIL import Image
|
| try:
|
| from sklearn.cluster import KMeans
|
| except ImportError:
|
| if method == 'KMeans':
|
| print('"sklearn" library is not installed, switching the extraction method to Median cut.')
|
| method = "Median cut"
|
|
|
| class Color:
|
| luminance_weights = np.array([0.2126, 0.7152, 0.0722])
|
|
|
| def __init__(self, RGB, frequency):
|
| self.rgb = tuple([c for c in RGB])
|
| self.freq = frequency
|
|
|
| def display(self, w=50, h=50):
|
| """
|
| Displays the represented color in a w x h window.
|
| :param w: width in pixels
|
| :param h: height in pixels
|
| """
|
|
|
| img = Image.new("RGB", size=(w, h), color=self.rgb)
|
| img.show()
|
|
|
| def __lt__(self, other):
|
| return self.freq < other.freq
|
|
|
| def get_colors(self, colorspace="rgb"):
|
| """
|
| Get the color in terms of a colorspace (string).
|
| :param colorspace: rgb/hsv/hls
|
| :return: corresponding color values
|
| """
|
| colors = {"rgb": self.rgb, "hsv": self.hsv, "hls": self.hls}
|
| return colors[colorspace]
|
|
|
| @property
|
| def hsv(self):
|
| return colorsys.rgb_to_hsv(*self.rgb)
|
|
|
| @property
|
| def hls(self):
|
| return colorsys.rgb_to_hls(*self.rgb)
|
|
|
| @property
|
| def luminance(self):
|
| return np.dot(self.luminance_weights, self.rgb)
|
|
|
| class ColorBox:
|
| """
|
| Represents a box in the RGB color space, with associated attributes, used in the Median Cut algorithm.
|
| """
|
| def __init__(self, colors):
|
| """
|
| Initialize with a numpy array of RGB colors.
|
| :param colors: np.ndarray (width * height, 3)
|
| """
|
|
|
| self.colors = colors
|
| self._get_min_max()
|
|
|
| def _get_min_max(self):
|
| min_channel = np.min(self.colors, axis=0)
|
| max_channel = np.max(self.colors, axis=0)
|
|
|
| self.min_channel = min_channel
|
| self.max_channel = max_channel
|
|
|
| def __lt__(self, other):
|
| """
|
| Compare cubes by volume
|
| :param other:
|
| """
|
| return self.size < other.size
|
|
|
| @property
|
| def size(self):
|
| return self.volume
|
|
|
| def _get_dominant_channel(self):
|
| dominant_channel = np.argmax(self.max_channel - self.min_channel)
|
| return dominant_channel
|
|
|
| @property
|
| def average(self):
|
| """
|
| Returns the average color contained in ColorBox
|
| :return: [R, G, B]
|
| """
|
|
|
| return np.mean(self.colors, axis=0)
|
|
|
| @property
|
| def volume(self):
|
| return np.prod(
|
| self.max_channel - self.min_channel,
|
| )
|
|
|
| def split(self):
|
| """
|
| Splits the ColorBox into two ColorBoxes at the median of the dominant color channel.
|
| :return: [ColorBox1, ColorBox2]
|
| """
|
|
|
|
|
| dominant_channel = self._get_dominant_channel()
|
|
|
|
|
| self.colors = self.colors[self.colors[:, dominant_channel].argsort()]
|
|
|
| median_index = len(self.colors) // 2
|
|
|
| return [
|
| ColorBox(self.colors[:median_index]),
|
| ColorBox(self.colors[median_index:]),
|
| ]
|
|
|
| class Palette:
|
| def __init__(self, colors):
|
| """
|
| Initializes a color palette with a list of Color objects.
|
| :param colors: a list of Color-objects
|
| """
|
|
|
| self.colors = colors
|
| self.frequencies = [c.freq for c in colors]
|
| self.number_of_colors = len(colors)
|
|
|
| def get_image(self, w=50, h=50):
|
| img = Image.new("RGB", size=(w * self.number_of_colors, h))
|
| arr = np.asarray(img).copy()
|
| for i in range(self.number_of_colors):
|
| c = self.colors[i]
|
| arr[:, i * h : (i + 1) * h, :] = c.rgb
|
| img = Image.fromarray(arr, "RGB")
|
| return img
|
|
|
| def k_means_extraction(arr, height, width, palette_size):
|
| """
|
| Extracts a color palette using KMeans.
|
| :param arr: pixel array (height, width, 3)
|
| :param height: height
|
| :param width: width
|
| :param palette_size: number of colors
|
| :return: a palette of colors sorted by frequency
|
| """
|
| arr = np.reshape(arr, (width * height, -1))
|
| model = KMeans(n_clusters=palette_size)
|
| labels = model.fit_predict(arr)
|
| palette = np.array(model.cluster_centers_, dtype=int)
|
| color_count = np.bincount(labels)
|
| color_frequency = color_count / float(np.sum(color_count))
|
| colors = []
|
| for color, freq in zip(palette, color_frequency):
|
| colors.append(Color(color, freq))
|
| return colors
|
|
|
| def median_cut_extraction(arr, height, width, palette_size):
|
| """
|
| Extracts a color palette using the median cut algorithm.
|
| :param arr:
|
| :param height:
|
| :param width:
|
| :param palette_size:
|
| :return:
|
| """
|
| arr = arr.reshape((width * height, -1))
|
| c = [ColorBox(arr)]
|
| full_box_size = c[0].size
|
|
|
| while len(c) < palette_size:
|
| largest_c_idx = np.argmax(c)
|
|
|
| c = c[:largest_c_idx] + c[largest_c_idx].split() + c[largest_c_idx + 1 :]
|
| colors = [Color(map(int, box.average), box.size / full_box_size) for box in c]
|
| return colors
|
|
|
| sort_methods = {
|
| "luminance": lambda c: c.luminance,
|
| "hue": lambda c: c.hsv[0],
|
| "saturation": lambda c: c.hsv[1],
|
| "value": lambda c: c.hsv[2],
|
| "lightness": lambda c: c.hls[2],
|
| }
|
|
|
| def extract_colors(image, palette_size=5, resize=True, mode="Median cut", sort_mode=None):
|
| """
|
| Extracts a set of 'palette_size' colors from the given image.
|
| :param image: PIL.Image object of path to Image file
|
| :param palette_size: number of colors to extract
|
| :param resize: whether to resize the image before processing, yielding faster results with lower quality
|
| :param mode: the color quantization algorithm to use. Currently supports K-Means (KM) and Median Cut (MC)
|
| :param sort_mode: sort colors by luminance, or by frequency
|
| :return: a list of the extracted colors
|
| """
|
| if isinstance(image, Image.Image):
|
| img = image
|
| else:
|
| img = Image.open(image)
|
| img = img.convert("RGB")
|
| if resize:
|
| img = img.resize((256, 256))
|
| width, height = img.size
|
| arr = np.asarray(img)
|
|
|
| if mode == "KMeans":
|
| colors = k_means_extraction(arr, height, width, palette_size)
|
| elif mode == "Median cut":
|
| colors = median_cut_extraction(arr, height, width, palette_size)
|
| else:
|
| raise NotImplementedError("Extraction mode not implemented!")
|
|
|
| if sort_mode in sort_methods:
|
| colors.sort(key=sort_methods.get(sort_mode), reverse=False)
|
| else:
|
| raise NotImplementedError("Sorting mode not implemented!")
|
| return Palette(colors)
|
|
|
|
|
| if(not overwrite):
|
| basename = f"_palette_{palette_size}x"
|
| else:
|
| p.do_not_save_samples = True
|
|
|
| proc = process_images(p)
|
|
|
|
|
|
|
| if len(proc.images) > 1:
|
| iter_offset = 1
|
| iter_num = len(proc.images) - 1
|
| else:
|
| iter_offset = 0
|
| iter_num = 1
|
|
|
| for i in range(iter_num):
|
| pal = extract_colors(proc.images[i+iter_offset], palette_size=palette_size, sort_mode=sort_by, mode=method)
|
| proc.images[i+iter_offset] = pal.get_image()
|
|
|
| images.save_image(proc.images[i+iter_offset], p.outpath_samples, basename,
|
| proc.seed + i, proc.prompt, opts.samples_format, info= proc.info, p=p)
|
|
|
| return proc |
