import gradio as gr
import os.path
from PIL import Image

def output(image_path,img_id,rule):
    print(image_path)
    output_path = './output'
    caption = 'person . mask'
    flag = False
    save_image_path = os.path.join(output_path,'test'+img_id+'.jpg')
    print(save_image_path)
    save_image = Image.open(save_image_path).convert("RGB")
    return save_image

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>服务监管效果Demo</center></h1>")
    gr.Markdown("""
    Tips:
    - 由于huggingface space不支持提供GPU,本Demo只支持cache中的三个样例效果展示
    - 目前基于GLIP的架构仅支持口罩佩戴监管,后续可以增加和选择其他规则
    - 判断部分采用的模型为ALBEF_14M.pth
    """)
    with gr.Row():
        with gr.Column():
            gr.Markdown("<h2>数据导入:</h2>")
            img_obj = gr.Image(value ='image/test.jpg',type = "filepath",label='object_img')
            index = gr.Radio(["0", "1", "2"],value ='0',label='img_id')
            gr.Markdown("<h2>规则选择(目前仅支持口罩佩戴规则):</h2>")
            rule = gr.Dropdown(["口罩佩戴规则", "隔位就坐规则", "核酸检测规则"],value='请选择规则...',label='规则选择')
        with gr.Column():
            gr.Markdown("<h2>效果展示:</h2>")
            img_output = gr.Image(type = "pil",label='output_img')
            btn = gr.Button("Submit",variant="primary")
    
    with  gr.Row():
        gr.Examples([['image/test.jpg','0','口罩佩戴规则'],['image/test1.jpg','1','口罩佩戴规则'],['image/test2.jpg','2','口罩佩戴规则']],[img_obj,index,rule],[img_output],output,True)

    btn.click(fn=output,inputs=[img_obj,index,rule],outputs=[img_output],)

demo.launch()