# -*- 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(
"",
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",
],
"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)