import gradio as gr

from algorithm import SearchImageTask

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # 作品登记证书搜索

    ## 使用方法
    1. 上传登记证书PDF文件，构建数据库
        - 一次可以上传一份或多份PDF文件
        - 每份PDF文件可以包含一份作品，也可以是多份作品
        - 一份作品的登记证书必须包含两页内容：第一页为作品信息，第二页为作品图片
    2. 上传待搜索图片，匹配作品登记证书
        - 每张图片会返回top-2的匹配结果

    > 多次上传数据，需要在前一次上传后点击右上角的'x'，然后再次上传数据

    ## 欢迎使用
    """
    )
    with gr.Row():
        with gr.Column():
            input_pdf = gr.Files(label="Upload PDF", height=300)
        with gr.Column():
            gallery_pdf = gr.Gallery(label="Product Images", elem_id="gallery", height=300)
    with gr.Row():
        upload_pdf_btn = gr.Button("上传PDF文件，构建数据库")

    upload_pdf_btn.click(SearchImageTask.upload_products_pdf_file, input_pdf, gallery_pdf)
    with gr.Row():
        with gr.Column(scale=1):
            gallery_input_image = gr.Gallery(label="Upload Images", elem_id="gallery", height=400)
        with gr.Column(scale=1):
            gallery_output_image = gr.Gallery(label="Matched Images", elem_id="gallery", height=400)

    with gr.Row():
        upload_img_btn = gr.Button("上传测试图片，匹配作品登记证书")

    upload_img_btn.click(SearchImageTask.upload_wait2search_image, gallery_input_image, gallery_output_image)


if __name__ == "__main__":
    demo.launch()