import gradio as gr
from rapid_unwrap.inference import DocUnwrapper

# 初始化模型
model_path = "models/uvdoc.onnx"
doc_unwrapper = DocUnwrapper(model_path=model_path)
# 添加示例
example_images = [
    "images/demo1.jpg",
    "images/demo1.png",
    "images/demo2.png",
    "images/demo3.jpg",

]


def process_image(img_path):
    unwrapped_img, elapse = doc_unwrapper(img_path)
    print(f"doc unwrap elapse: {elapse}")
    return unwrapped_img


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'>RapidUnWrap</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
                )
                run_button = gr.Button("summit")
            with gr.Column(scale=2):  # 右边占2/3
                output_image = gr.Image(label="output")

        # 绑定按钮点击事件
        run_button.click(fn=process_image, inputs=img_input, outputs=output_image)

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