Feature: Color Palette support

[chucktobbes]

Hi @Astropulse, I tried to add to restrict the reduced colors to a color palette. But all I got is a wild mix of colors. Do you have any idea why this happens and how to improve it? Here is the result, the palette and my code: Used Color Palette: endesga-32.txt

import modules.scripts as scripts
import hitherdither
import gradio as gr

import cv2
import numpy as np
from PIL import Image

from modules import images
from modules.processing import Processed, process_images
from modules.shared import opts, cmd_opts, state

from torch import Tensor
from torch.nn import Conv2d
from torch.nn import functional as F
from torch.nn.modules.utils import _pair
from typing import Optional

script_dir = scripts.basedir()

# Runs cv2 k_means quantization on the provided image with "k" color indexes

def palettize(img, k, d):
    image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    image = Image.fromarray(image).convert("RGB")

    if d > 0:
        if k <= 64:
            img_indexed = image.quantize(
                colors=k, method=1, kmeans=k, dither=0).convert('RGB')

            palette = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] # red, green, blue
            palette = hitherdither.palette.Palette(palette)
            img_indexed = hitherdither.ordered.yliluoma.yliluomas_1_ordered_dithering(
            image, palette, order=2**d).convert('RGB')
        else:
            img_indexed = image.quantize(
                colors=k, method=1, kmeans=k, dither=0).convert('RGB')
    else:
        img_indexed = image.quantize(
            colors=k, method=1, kmeans=k, dither=0).convert('RGB')

    result = cv2.cvtColor(np.asarray(img_indexed), cv2.COLOR_RGB2BGR)
    return result

def hex_to_rgb(hex_color):
    hex_color = hex_color[-6:]
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

class Script(scripts.Script):
    def title(self):
        return "Palettize"

    def show(self, is_img2img):
        return True

    def ui(self, is_img2img):
        clusters = gr.Slider(minimum=2, maximum=128, step=1,
                             label='Colors in palette', value=24)
        downscale = gr.Checkbox(
            label='Downscale before processing', value=True)
        with gr.Row():
            scale = gr.Slider(minimum=2, maximum=32, step=1,
                              label='Downscale factor', value=8)
            dither = gr.Slider(minimum=0, maximum=3, step=1,
                               label='Dithering', value=0)
        file = gr.File(label='Chosse a .txt file with hex colors on per line')

        return [downscale, scale, clusters, dither, file]

    def run(self, p, downscale, scale, clusters, dither, file):

        if dither > 0:
            if clusters <= 64:
                print(
                    f'Palettizing output to {clusters} colors with order {2**dither} dithering...')
            else:
                print('Palette too large, max colors for dithering is 64.')
                print(f'Palettizing output to {clusters} colors...')
        else:
            print(f'Palettizing output to {clusters} colors...')

        processed = process_images(p)

        generations = p.batch_size*p.n_iter

        for i in range(generations + int(generations > 1)):
            # Converts image from "Image" type to numpy array for cv2
            img = np.array(processed.images[i]).astype(np.uint8)

            if downscale:
                img = cv2.resize(img, (int(
                    img.shape[1]/scale), int(img.shape[0]/scale)), interpolation=cv2.INTER_LINEAR)

             # Map colors to nearest palette color
            if file is not None:
                with open(file.name, 'r', encoding='utf-8') as f:
                    content = f.read()
                    if content:
                        palette = []
                        palette = [hex_to_rgb(line.strip()) for line in content.splitlines()]
                        print(f'Loaded colors: {palette}')       
                        clusters = len(palette)
                        img = palettize(img, clusters, dither)
                        # Convert the img NumPy array to a PIL Image object
                        img = Image.fromarray(img)
                        # Quantize the image to reduce the number of colors
                        img = img.quantize(colors=len(palette))
                        # Map the colors in the quantized image to the colors in the palette
                        recolor_image = Image.new('RGB', img.size)
                        recolor_image.putdata([palette[color] for color in img.getdata()])
                        # save the recolored image
                        img = recolor_image
                        # Convert the PIL Image object to a NumPy array
                        img = np.array(img)
            else:
                img = palettize(img, clusters, dither)

            if downscale:
                img = cv2.resize(img, (int(
                    img.shape[1]*scale), int(img.shape[0]*scale)), interpolation=cv2.INTER_NEAREST)

            processed.images[i] = Image.fromarray(img)
            images.save_image(processed.images[i], p.outpath_samples, "palettized", processed.seed +
                            i, processed.prompt, opts.samples_format, info=processed.info, p=p)

            if generations > 1:
                grid = images.image_grid(
                    processed.images[1:generations+1], p.batch_size)
                processed.images[0] = grid

            if opts.grid_save:
                images.save_image(processed.images[0], p.outpath_grids, "palettized",
                                prompt=p.prompt, seed=processed.seed, grid=True, p=p)

        return processed

[Astropulse]

Oh wow. Looks like the image is getting indexed, but the indexes are in the wrong order. There is some structure to the output, its just all mixed up.

I think part of the problem is the .txt file parser putting things in a weird order.

I'll put together an official version of this. Here I have a script that already does something similar, I just need to port it over to this extension.

[chucktobbes]

That would be very nice, because my ideas of solving the problem start to fade. Even with different attempts and support of GPT4 I only get the color mix.

[Astropulse]

Added the image palette feature.