import gradio as gr
import os
from PIL import Image
from torchkeras import plots 
from torchkeras.data import get_url_img
from typing import List

# 定义需要保存上传图像的文件夹路径
UPLOAD_FOLDER = "app_src/input"
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

# 定义需要显示的图像文件路径
DISPLAY_IMAGE_PATH = "app_src/output/test_1_seg_result.png"

# 输入图像和文本的接口
def input_fn(img, text):
    if isinstance(img,str):
        try:
            img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
        except:
            raise gr.Error("Failed to get image from URL, please upload from local.")

    image_path = os.path.join(UPLOAD_FOLDER, 'input.png')
    img.save(image_path)
    output_data = output_fn(text)
    return output_data

# 输出图像的接口
def output_fn(text):
    # 返回指定路径的图像
    img = Image.open(DISPLAY_IMAGE_PATH).convert('RGB')
    info = {"size": img.size, "format": img.format}
    text = "<center><h1><span style='color:green'>Successful detection! 🎉</span></h1></center>"
    return img, text


if __name__ == "__main__":
    # 创建 Gradio 接口
    flag=0
    with gr.Blocks() as demo:
        gr.Markdown("# TOIST v2")
        with gr.Row():
            with gr.Tab("Camera"):
                with gr.Row():
                    with gr.Column():
                        camera = gr.Image(source='webcam',type='pil').style( height=300)
                        text = gr.inputs.Textbox(label="Please enter prompt, what do you want to do ?"),  # 文本输入组件
                        camera_button = gr.Button("Submit",variant="primary")

            with gr.Tab("Url"):
                with gr.Row():
                    with gr.Column():
                        default_url = ''
                        url = gr.Textbox(value=default_url)
                        text = gr.Textbox(label="Please enter prompt, what do you want to do ?")  # 文本输入组件
                        url_button = gr.Button("Submit",variant="primary")
                        # examples = gr.Examples = [UPLOAD_FOLDER+'exp1.png',UPLOAD_FOLDER+'exp2.png']
            
            with gr.Tab("Upload"):
                with gr.Row():
                    with gr.Column(equal_weight=True):
                        input_img = gr.Image(type='pil').style( height=300)
                        text = gr.Textbox(label="Please enter prompt, what do you want to do ?")  # 文本输入组件
                        img_button = gr.Button("Submit",variant="primary")
                        # examples = gr.Examples(examples=[UPLOAD_FOLDER+'exp1.png',UPLOAD_FOLDER+'exp2.png'])
                                               
            
            with gr.Column(equal_weight=True):
                gr.Markdown("<h3> Output </h3>")
                output = [
                            gr.Image(type='pil', label='Image').style(height=400, align="center"),
                            # gr.Textbox(label="Successful detection!<span style='color: green'>&#x1F600;</span><br> 🎉  \nYour prompt: ", type='text'),
                            gr.Markdown()
                            # gr.JSON(type='text')
                            ]
                
            camera_button.click(input_fn,
                        [camera,text], 
                        output
                        )
            url_button.click(input_fn,
                        [url,text], 
                        output
                        )
            img_button.click(input_fn,
                        [input_img,text], 
                        output
                        )
        examples = gr.Examples([[UPLOAD_FOLDER+'/exp1.png','sit on something'],
                                [UPLOAD_FOLDER+'/exp2.png','sit on something']],
                                               inputs=[input_img,text],
                                               outputs=output,
                                               fn = input_fn)

    gr.close_all() 
    demo.queue()
    demo.launch()
        
    # gr.Interface(
    #     fn=input_fn,  # 输出函数
    #     inputs=[
    #         gr.inputs.Image(label="upload your image", type="pil" ,input_size=(256, 256)),  # 图像输入组件
    #         gr.inputs.Textbox(label="Please enter prompt, what do you want to do ?"),  # 文本输入组件
    #     ],
    #     outputs=gr.outputs.Image(type="pil",image_size=(256, 256)),  # 图像输出组件
    #     title="TOIST v2",
    #     description="Experience TOIST v2",
    #     theme="github",  # 采用 GitHub 风格
    #     server_name="0.0.0.0",  # 指定服务器地址
    #     server_port=7878,  # 指定服务器端口
    #     # auth=gr.auth.LoginRequired(authenticate),  # 需要登录验证
    # ).launch()