# -*- 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(
"",
unsafe_allow_html=True,
)
st.markdown(
"""
""",
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",
],
}
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)