sd-palettize

sd-palettize.py

#!/usr/bin/env python
"""\
This script is from Astropulse, Written by ChatGPT-4o into a ready-to-use script version.
<https://github.com/Astropulse/sd-palettize>

---

该脚本，来自 Astropulse，经过 ChatGPT-4o 编写为了可直接使用的脚本版本。
<https://github.com/Astropulse/sd-palettize>

Usage: python sd-palettize.py path/to/pixle/art.png
"""
import sys
import numpy as np
from PIL import Image
from sklearn.cluster import KMeans
import os

def extract_palette(image, num_colors):
    # 将图像转换为 RGB
    image = image.convert('RGB')
    # 获取图像像素数据
    pixels = np.array(image)
    pixels = pixels.reshape((-1, 3))

    # 使用 k-means 聚类算法提取调色板
    kmeans = KMeans(n_clusters=num_colors, random_state=0).fit(pixels)
    palette = kmeans.cluster_centers_.astype(int)

    return palette

def save_palette(palette, file_path):
    # 保存调色板为 .gpl 文件
    with open(file_path, 'w') as f:
        f.write("GIMP Palette\n")
        f.write("Name: Palette\n")
        f.write("Columns: 16\n")
        f.write("#\n")
        for color in palette:
            f.write(f"{color[0]} {color[1]} {color[2]} Untitled\n")

def quantize_image(image, palette):
    # 将调色板应用于图像
    palette_image = Image.new('P', (1, len(palette)))
    palette_image.putpalette(palette.flatten().tolist())
    image = image.convert('RGB').quantize(palette=palette_image, dither=Image.NONE)
    return image.convert('RGB')

def main(input_path):
    # 打开图像
    image = Image.open(input_path)

    # 提取调色板
    num_colors = 256  # 最大颜色数量，可以根据需要调整
    palette = extract_palette(image, num_colors)

    # 保存调色板
    palette_path = os.path.splitext(input_path)[0] + '.gpl'
    save_palette(palette, palette_path)

    # 量化图像
    quantized_image = quantize_image(image, palette)

    # 保存量化后的图像
    output_path = os.path.splitext(input_path)[0] + '_sd-palettize.png'
    quantized_image.save(output_path)

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: python sd-palettize.py [file-name]")
        sys.exit(1)

    input_path = sys.argv[1]
    main(input_path)