import gradio as gr
import numpy as np
from sklearn.cluster import KMeans
from sklearn.utils import shuffle
from sklearn.cluster import MeanShift as ms
from PIL import Image

模式=gr.inputs.CheckboxGroup(choices=['kmeans','meanshift'])
颜色数=gr.inputs.Slider(maximum=25,minimum=1,step=1)
宽度=gr.inputs.Number()
输入图片=gr.inputs.Image()

def func(输入图片,颜色数,宽度,模式):
    china = 输入图片
    n=int(颜色数)
    width=宽度
    mode=模式[0]
    china = np.array(china, dtype=np.float64)/china.max()
    w, h, d = original_shape = tuple(china.shape)
    image_array = np.reshape(china, (w*h, d))
    image_array_sample = shuffle(image_array, random_state=0)[:1000]
    if(mode=='kmeans'):
        kmeans = KMeans(n_clusters=n, random_state=0).fit(image_array_sample)
    if(mode=='meanshift'):
        kmeans = ms(bandwidth=width).fit(image_array_sample)
    labels = kmeans.predict(image_array)
    image_kmeans = image_array.copy()
    for i in range(w*h):
        image_kmeans[i] = kmeans.cluster_centers_[labels[i]]
    image_kmeans = image_kmeans.reshape(w, h, d)
    return image_kmeans

title="图片颜色种类压缩（kmeans和meanshift）"
description="""
            颜色数指输出图片的颜色数量，为kmeans模式专属，建议取值2到20，越大运算越慢。
            宽度用于衡量输出图片的精细度，越小越精细，为meanshift模式专属，建议取值0.05到0.3，越小运算越慢。
    """

demo = gr.Interface(
    fn=func,
    inputs=[输入图片,颜色数,宽度,模式],
    outputs='image',
    title=title,
    description=description,
    live=True,
    flagging_callback=gr.SimpleCSVLogger()
)
demo.launch(share=True)