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