import cv2  
from PIL import Image, ImageEnhance
import gradio as gr
from sklearn.cluster import KMeans
import numpy as np

with gr.Blocks() as interface:

    with gr.Row():
        n_colors = gr.Slider(2, 32, 12, step=1, label="图片要加工的目标颜色数量")
     
    with gr.Row():
        img_input = gr.Image()
        img_output = gr.Image()
    
    section_btn1 = gr.Button("合并色彩")
    
    # 图片模型训练
    def img_fit_predict(img,n_colors):
        data = img.reshape(-1,3)
        # 把原始图片压缩成n_colors个颜色
        kmeans = KMeans(n_clusters=n_colors)
        y_ = kmeans.fit_predict(data)
        # 模型合并颜色
        colors = kmeans.cluster_centers_/255
        output_temp = colors[y_].reshape(img.shape)
        return output_temp
    
    section_btn1.click(img_fit_predict, inputs=[img_input,n_colors], outputs=img_output)
    
    with gr.Row():
        gaussian_blur = gr.Slider(1, 13, 13, step=2, label="整体降噪参数调整")
        structuring_element = gr.Slider(1, 13, 3, step=2, label="去除小噪声")
        canny_start = gr.Slider(1, 200, 4, step=1, label="边缘检测-开始参数")
        canny_end = gr.Slider(1, 200, 10, step=1, label="边缘检测-结束参数")
    
    with gr.Row():
        thresh_val = gr.Slider(50, 500, 205, step=1, label="二值图像-thresh")
        maxval = gr.Slider(50, 500, 330, step=1, label="二值图像-maxval")
        enhance = gr.Slider(0, 1, 0.8, step=0.1, label="增强颜色-enhance")
        blend = gr.Slider(0, 1, 0.4, step=0.1, label="增强颜色-blend")
     
    section_btn2 = gr.Button("调整图片")   
    with gr.Row():
        closed_output = gr.Image()
        img_param_output = gr.Image()

    # 调整模型结果参数
    def turn_arguments(img,img_output,gaussian_blur,structuring_element,canny_start,canny_end,thresh_val,maxval,enhance,blend):
        gray = cv2.cvtColor(img_output, cv2.COLOR_BGR2GRAY) 
        # 对灰度图像进行高斯滤波，以去除噪声  
        gray = cv2.GaussianBlur(gray, (gaussian_blur,gaussian_blur), 0)  
        # 使用Canny算子进行边缘检测  
        edges = cv2.Canny(gray, canny_start, canny_end)
        # 将边缘图像转换为二值图像  
        _, thresh = cv2.threshold(edges, thresh_val, maxval, cv2.THRESH_BINARY) 
        # 对二值图像进行形态学操作，以去除小的噪点  
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (structuring_element, structuring_element))  
        closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        image = Image.fromarray(img_output)
        closed = closed.astype(img.dtype)
        result = cv2.bitwise_and(img_output, img_output, mask=closed)
        result[closed==0] = (255,255,255)
        # 颜色空间转换  
        enhancer = ImageEnhance.Color(image=image)
        # 增强颜色  
        img1 = enhancer.enhance(enhance).convert('RGB')
        img2 = Image.fromarray(result).convert('RGB')
        union_img = np.asarray(Image.blend(img2, img1, blend))
        return result,union_img
    
    section_btn2.click(turn_arguments,inputs=[img_input, img_output,gaussian_blur,
                                              structuring_element,canny_start,canny_end,thresh_val,maxval,enhance,blend ],
                                    outputs = [closed_output,img_param_output])

interface.launch()