import gradio as gr
from colorthief import ColorThief
import numpy as np
from io import BytesIO

def extract_colors(image, num_colors):
    img = BytesIO()
    image.save(img, format='PNG')
    img.seek(0)

    color_thief = ColorThief(img)
    palette = color_thief.get_palette(color_count=num_colors)
    colors = np.array(palette)

    return colors

def plot_colors(colors):
    palette = np.zeros((50, len(colors) * 50, 3), dtype=int)
    step = 50
    for idx, color in enumerate(colors):
        palette[:, idx*step:(idx+1)*step, :] = color
    return palette

def colors_to_hex(colors):
    hex_colors = ['#%02x%02x%02x' % tuple(color) for color in colors]
    return hex_colors

def process_image(image, num_colors):
    colors = extract_colors(image, num_colors)
    palette = plot_colors(colors)
    palette_str = "palette=" + str(colors.tolist())
    hex_colors = colors_to_hex(colors)
    hex_palette_str = 'palette=' + str(hex_colors)
    return palette, palette_str, hex_palette_str

def main():
    iface = gr.Interface(fn=process_image,
                         inputs=[
                             gr.Image(type="pil", label="上传图片"),
                             gr.Slider(minimum=1, maximum=20, value=5, step=1, label="调色板颜色数量")
                         ],
                         outputs=[
                             gr.Image(type="numpy", label="生成的调色板"),
                             gr.Textbox(label="调色板RGB值"),
                             gr.Textbox(label="调色板16进制值")
                         ])
    
    iface.launch()

if __name__ == '__main__':
    main()