File size: 2,786 Bytes
628d3b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9b5cb05
628d3b9
9b5cb05
 
 
 
628d3b9
 
 
 
 
 
 
 
 
 
78410a0
 
 
628d3b9
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
# -*- 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",
        "figure",
        "figure_caption",
        "table",
        "table_caption",
        "header",
        "footer",
        "reference",
        "equation",
    ],
    "yolov8n_layout_report": [
        "text",
        "title",
        "header",
        "footer",
        "figure",
        "figure_caption",
        "table",
        "table_caption",
        "toc",
    ],
}
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)