Update app.py

app.py

@@ -1,102 +1,99 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-import numpy as np
-import streamlit as st
-from PIL import Image
-import cv2
+import gradio as gr
 
-from rapid_layout import RapidLayout, VisLayout
+from rapid_layout import ModelType, RapidLayout, RapidLayoutInput
+
+
+def get_layout_res(img_input, model_type, conf_thresh, iou_thresh):
+    cfg = RapidLayoutInput(
+        model_type=ModelType(model_type), conf_thresh=conf_thresh, iou_thresh=iou_thresh
+    )
+    engine = RapidLayout(cfg=cfg)
+    result = engine(img_input)
+    vised_img = result.vis("layout_vis.jpg")
+    return vised_img, f"{result.elapse:.4f}"
 
-st.markdown(
-    "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidLayout' style='text-decoration: none'>Rapid Layout</a></h1>",
-    unsafe_allow_html=True,
-)
-st.markdown(
-    """
-<p align="left">
-    <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.13-aff.svg"></a>
-    <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
-    <a href="https://pypi.org/project/rapid-layout/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-layout"></a>
-    <a href="https://pepy.tech/project/rapid-layout"><img src="https://static.pepy.tech/personalized-badge/rapid-layout?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></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>
-</p>
-""",
-    unsafe_allow_html=True,
-)
 
-model_types = {
-    "pp_layout_cdla": [
-        "text",
-        "title",
-        "figure",
-        "figure_caption",
-        "table",
-        "table_caption",
-        "header",
-        "footer",
-        "reference",
-        "equation",
-    ],
-    "pp_layout_publaynet": ["text", "title", "list", "table", "figure"],
-    "pp_layout_table": ["table"],
-    "yolov8n_layout_paper": [
-        "Text",
-        "Title",
-        "Header",
-        "Footer",
-        "Figure",
-        "Table",
-        "Toc",
-        "Figure caption",
-        "Table caption",
-    ],
-    "yolov8n_layout_report": [
-        "Text",
-        "Title",
-        "Header",
-        "Footer",
-        "Figure",
-        "Table",
-        "Toc",
-        "Figure caption",
-        "Table caption",
-    ],
-    "yolov8n_layout_publaynet": ["Text", "Title", "List", "Table", "Figure"],
-    "yolov8n_layout_general6": [
-        "Text",
-        "Title",
-        "Figure",
-        "Table",
-        "Caption",
-        "Equation",
-    ],
-    "doclayout_docstructbench": ['title', 'plain text', 'abandon', 'figure', 'figure_caption', 'table', 'table_caption', 'table_footnote', 'isolate_formula', 'formula_caption'],
-    "doclayout_d4la": ['DocTitle', 'ParaTitle', 'ParaText', 'ListText', 'RegionTitle', 'Date', 'LetterHead', 'LetterDear', 'LetterSign', 'Question', 'OtherText', 'RegionKV', 'RegionList', 'Abstract', 'Author', 'TableName', 'Table', 'Figure', 'FigureName', 'Equation', 'Reference', 'Footer', 'PageHeader', 'PageFooter', 'Number', 'Catalog', 'PageNumber'],
-    "doclayout_docsynth": ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']
-}
-select_model = st.selectbox("选择版面分析模型：", model_types.keys())
-st.write("支持检测类型：")
-st.code(model_types[select_model], language="python")
+custom_css = """
+    body {font-family: body {font-family: 'Helvetica Neue', Helvetica;}
+    .gr-button {background-color: #4CAF50; color: white; border: none; padding: 10px 20px; border-radius: 5px;}
+    .gr-button:hover {background-color: #45a049;}
+    .gr-textbox {margin-bottom: 15px;}
+    .example-button {background-color: #1E90FF; color: white; border: none; padding: 8px 15px; border-radius: 5px; margin: 5px;}
+    .example-button:hover {background-color: #FF4500;}
+    .tall-radio .gr-radio-item {padding: 15px 0; min-height: 50px; display: flex; align-items: center;}
+    .tall-radio label {font-size: 16px;}
+    .output-image, .input-image, .image-preview {height: 300px !important}
+"""
+
+with gr.Blocks(title="Rapid Layout", css="custom_css", theme=gr.themes.Soft()) as demo:
+    gr.HTML(
+        """
+        <h1 style='text-align: center;font-size:40px'>Rapid Layout v1</h1>
+
+        <div style="display: flex; justify-content: center; gap: 10px;">
+            <a><img src="https://img.shields.io/badge/DeployVersion-v1.0.0-brightgree"></a>
+            <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.13-aff.svg"></a>
+            <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
+            <a href="https://pypi.org/project/rapid-layout/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-layout"></a>
+            <a href="https://pepy.tech/project/rapid-layout"><img src="https://static.pepy.tech/personalized-badge/rapid-layout?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></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>
+        </div>
+    """
+    )
+    with gr.Row():
+        model_type = gr.Dropdown(
+            choices=[v.value for v in ModelType],
+            label="版面分析模型",
+            value=ModelType.PP_LAYOUT_CDLA.value,
+            interactive=True,
+        )
+        conf_thresh = gr.Slider(
+            label="conf_thresh",
+            minimum=0,
+            maximum=1.0,
+            value=0.5,
+            step=0.1,
+            info="置信度",
+            interactive=True,
+        )
+        iou_thresh = gr.Slider(
+            label="IoU Threshold",
+            minimum=0,
+            maximum=1.0,
+            value=0.5,
+            step=0.1,
+            info="IoU阈值设定",
+            interactive=True,
+        )
 
-conf_threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5)
-iou_threshold = st.slider("IoU Threshold", 0.0, 1.0, 0.5)
+    with gr.Row():
+        run_btn = gr.Button("Run")
 
-with st.spinner(f"Downloading {select_model} model..."):
-    layout_engine = RapidLayout(model_type=select_model, conf_thres=conf_threshold, iou_thres=iou_threshold)
+    with gr.Row():
+        img_input = gr.Image(label="Upload or Select Image", sources="upload")
+        img_output = gr.Image(label="Output Image", show_download_button=True)
+    elapse = gr.Textbox(label="Elapse(s)")
 
-st.write("请上传图像：")
-img_suffix = ["png", "jpg", "jpeg"]
-img_file_buffer = st.file_uploader(
-    "Upload an image", type=img_suffix, key="layout", label_visibility="collapsed"
-)
+    run_btn.click(
+        get_layout_res,
+        inputs=[img_input, model_type, conf_thresh, iou_thresh],
+        outputs=[img_output, elapse],
+    )
 
-if img_file_buffer:
-    image = Image.open(img_file_buffer)
-    img = np.array(image)
-    with st.spinner("推理中...."):
-        boxes, scores, class_names, *elapse = layout_engine(img, )
-        ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
+    examples = gr.Examples(
+        examples=[
+            ["images/layout.jpg", ModelType.PP_LAYOUT_CDLA.value, 0.5, 0.5],
+        ],
+        examples_per_page=5,
+        inputs=[img_input, model_type, conf_thresh, iou_thresh],
+        fn=get_layout_res,
+        outputs=[img_output, elapse],
+        cache_examples=False,
+    )
 
-    st.image(ploted_img, use_column_width=True)
+if __name__ == "__main__":
+    demo.launch(debug=True)