File size: 3,867 Bytes
628d3b9
 
 
 
 
 
e7d8bed
628d3b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9c64580
 
 
 
 
 
 
 
 
628d3b9
 
9c64580
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
628d3b9
31aae41
 
 
628d3b9
9b5cb05
628d3b9
9b5cb05
 
6482965
 
e7d8bed
9b5cb05
e7d8bed
628d3b9
 
 
 
 
 
 
 
 
 
78410a0
e7d8bed
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import numpy as np
import streamlit as st
from PIL import Image
import cv2

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_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")

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 st.spinner(f"Downloading {select_model} model..."):
    layout_engine = RapidLayout(model_type=select_model, conf_thres=conf_threshold, iou_thres=iou_threshold)

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)