from sqlite3 import InterfaceError
import gradio as gr
import cv2
from color_editor import process_function

examples = [
    [
        "6 33 58",
        "227 239 251",
        "",
        "0 30",
        "240 255",
        "2",
        "",
        "./resources/1.png",
        "./resources/mask1.png",
    ],
    [
        "21 21 40",
        "41 42 42",
        "",
        "5 10",
        "0 46",
        "4",
        "",
        "./resources/2.png",
        "./resources/mask2.png",
    ],
]

def str2lst(strr):
    if strr == '':
        return None

    item = strr.split(' ')
    lst = []
    for i in item:
        lst.append(int(i))
    return lst

def inference(src_rgb, tag_rgb, flag_h, flag_s, flag_v, v_contrast, curve, image, mask):
    # 整理输入参数
    cfg = {}
    cfg['img'] = image
    cfg['mask'] = mask
    cfg['src_rgb'] = str2lst(src_rgb)
    cfg['tag_rgb'] = str2lst(tag_rgb)

    cfg['flag_h'] = str2lst(flag_h)
    cfg['flag_s'] = str2lst(flag_s)
    cfg['flag_v'] = str2lst(flag_v)

    if v_contrast == '':
        cfg['v_contrast'] = 0
    else:
        cfg['v_contrast'] = int(v_contrast)
    cfg['dilate_kernel'] = 0

    # 打光参数暂时置为False
    cfg['if_light'] = False
    cfg['light_type'] = 2
    cfg['light_strength'] = 40

    
    if curve == "":
        cfg['if_curve'] = False
        cfg['l_rnk'] = 0
        cfg['h_rnk'] = 0
    else:
        curve_item = str2lst(curve)
        print('curve item : ', curve_item[0], curve_item[1])
        cfg['if_curve'] = True
        cfg['l_rnk'] = curve_item[0]
        cfg['h_rnk'] = curve_item[1]

    result = process_function(cfg)


    tag_color = result['tag_color']
    tag_hsv = result['tag_hsv']
    tag_recommand = result['tag_recommand']
    res_img = result['res_img']

    return tag_color, tag_hsv, tag_recommand, res_img

# 界面变量定义
title="🚗汽车改色"
description = """
<font size=4>
汽车改色使用说明: \n
将rgb色彩空间映射hsv(色调(**H**ue)-饱和度(**S**aturation)-亮度(**V**alue))色彩空间, 通过调整车漆部分hsv值来达到换色的效果, 点击底部exmpales进行快速尝试 \n
[[使用文档](https://bytedance.feishu.cn/docx/J1rbdzJbao2WKAxKuc0cOkhTn0e?from=from_copylink)]
</font>
"""

iface = gr.Interface(
    fn=inference, 
    inputs=[
        gr.Textbox(lines=1, placeholder=None, label="原图rgb｜eg. 255 255 255"),
        gr.Textbox(lines=1, placeholder=None, label="目标rgb｜eg. 255 255 255"),
        gr.Textbox(lines=1, placeholder=None, label="色调范围｜eg. 0 180"),
        gr.Textbox(lines=1, placeholder=None, label="饱和度范围｜eg. 0 180"),
        gr.Textbox(lines=1, placeholder=None, label="亮度范围｜eg. 0 180"),
        gr.Textbox(lines=1, placeholder=None, label="亮度拉伸系数｜选填｜eg. 1"),
        gr.Textbox(lines=1, placeholder=None, label="曲线工具｜选填｜eg. 20 240"),
        gr.Image(type="filepath", label="改色原图"),
        gr.Image(type="filepath", label="车漆mask"),
    ],
    outputs=[
        gr.Textbox(lines=1, placeholder=None, label="目标色属于"),
        gr.Textbox(lines=1, placeholder=None, label="目标rgb转hsv"),
        gr.Textbox(lines=1, placeholder=None, label="目标颜色参数范围｜供参考"),
        gr.Image(type="filepath", label="换色结果"),
    ],
    title=title,
    description=description,
    examples=examples,
)
iface.launch()