import gradio as gr
from rapid_undistorted.inference import InferenceEngine
import numpy as np
# 初始化模型
engine = InferenceEngine("config.yaml")
# 添加示例
example_images = [
    "images/demo.jpg",
    "images/demo1.jpg",
    "images/demo1.png",
    "images/demo2.png",
    "images/demo3.jpg",

]


# 定义任务和模型选项
tasks = {
    "unwrap": ["UVDoc", None],
    "unshadow": ["GCDnet",  None],
    "unblur": ["OpenCvBilateral", "NAFDPM", None],
    "binarize": ["UnetCnn", None]
}

def process_image(img_path, unwrap_model, unshadow_model, binarize_model, unblur_model):
    task_list = []
    if unwrap_model:
        task_list.append(("unwrap", unwrap_model))
    if unshadow_model:
        task_list.append(("unshadow", unshadow_model))
    if binarize_model:
        task_list.append(("binarize", binarize_model))
    if unblur_model:
        task_list.append(("unblur", unblur_model))

    unwrapped_img, elapse = engine(img_path, task_list)
    print(f"doc unwrap elapse: {elapse}")
    return unwrapped_img.astype(np.uint8),elapse


def main():
    # 定义Gradio界面
    with gr.Blocks(css="""
        .scrollable-container {
            overflow-x: auto;
            white-space: nowrap;
        }
        .header-links {
            text-align: center;
        }
        .header-links a {
            display: inline-block;
            text-align: center;
            margin-right: 10px;  /* 调整间距 */
        }
    """) as demo:
        gr.HTML(
            "<h1 style='text-align: center;'><a href='https://github.com/Joker1212/RapidUnWrap'>RapidUnDistort</a></h1>"
        )
        gr.HTML('''
                                            <div class="header-links">
                                              <a href=""><img src="https://img.shields.io/badge/Python->=3.8,<3.13-aff.svg"></a>
  <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Mac%2C%20Win-pink.svg"></a>
<a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
  <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
  <a href="https://github.com/RapidAI/TableStructureRec/blob/c41bbd23898cb27a957ed962b0ffee3c74dfeff1/LICENSE"><img alt="GitHub" src="https://img.shields.io/badge/license-Apache 2.0-blue"></a>
                                            </div>
                                            ''')
        with gr.Row():
            with gr.Column(scale=1):  # 左边占1/3
                img_input = gr.Image(label="Upload or Select Image", sources="upload", value="images/demo1.jpg")
                # 示例图片选择器
                examples = gr.Examples(
                    examples=example_images,
                    examples_per_page=len(example_images),
                    inputs=img_input,
                    fn=lambda x: x,  # 简单返回图片路径
                    outputs=img_input,
                    cache_examples=False
                )

                unwrap_dropdown = gr.Dropdown(label="Select Unwrap Model", choices=tasks["unwrap"], value="UVDoc")
                unshadow_dropdown = gr.Dropdown(label="Select Unshadow Model", choices=tasks["unshadow"], value="GCDnet")
                binarize_dropdown = gr.Dropdown(label="Select Binarize Model", choices=tasks["binarize"], value=None)
                unblur_dropdown = gr.Dropdown(label="Select Unblur Model", choices=tasks["unblur"], value="OpenCvBilateral")

                run_button = gr.Button("summit")

            with gr.Column(scale=2):  # 右边占2/3
                output_image = gr.Image(label="output")
                elapse_textbox = gr.Textbox(label="Elapsed Time", interactive=False)
        # 绑定按钮点击事件
        run_button.click(fn=process_image, inputs=[img_input, unwrap_dropdown, unshadow_dropdown, binarize_dropdown, unblur_dropdown],outputs=[output_image, elapse_textbox])

    # 启动应用
    demo.launch()
if __name__ == '__main__':
    main()