app.py

# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import numpy as np
import streamlit as st
from PIL import Image

from rapid_layout import RapidLayout, VisLayout

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_yolo": [
        "title",
        "text",
        "abandon",
        "figure",
        "figure_caption",
        "table",
        "table_caption",
        "table_footnote",
        "isolate_formula",
        "formula_caption",
    ],
}
select_model = st.selectbox("选择版面分析模型：", model_types.keys())
st.write("支持检测类型：")
st.code(model_types[select_model], language="python")

with st.spinner(f"Downloading {select_model} model..."):
    layout_engine = RapidLayout(model_type=select_model)

st.write("请上传图像：")
img_suffix = ["png", "jpg", "jpeg"]
img_file_buffer = st.file_uploader(
    "Upload an image", type=img_suffix, key="layout", label_visibility="collapsed"
)

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)

    st.image(ploted_img, use_column_width=True)