Spaces:
Running
Running
File size: 3,234 Bytes
628d3b9 9c64580 628d3b9 9c64580 628d3b9 9408e2f 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
# -*- 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)
|