# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import copy
import cv2
import numpy as np
import streamlit as st
from PIL import Image
from rapid_layout import RapidLayout
from rapid_orientation import RapidOrientation
from rapid_table import RapidTable
from streamlit_image_select import image_select
orientation_engine = RapidOrientation()
layout_engine = RapidLayout()
table_engine = RapidTable()
def orientation_part():
st.markdown('##### 文档图像方向分类')
img = image_select(
label="示例图（点击选择）",
images=[
'images/orientation/img_rot0_demo.jpg',
'images/orientation/rot_90.jpg',
'images/orientation/img_rot180_demo.jpg',
'images/orientation/rot_270.jpg'
],
captions=['0', '90', '180', '270'],
key='cls'
)
img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
key='orientation',
label_visibility='collapsed')
col1, col2 = st.columns([5, 5])
img_empty = col1.empty()
image, orientation_res, elapse = get_cls_res(img)
if img_file_buffer:
image, orientation_res, elapse = get_cls_res(img_file_buffer)
img_empty.image(image, use_column_width=True)
col2.markdown(f'- 方向分类结果：{orientation_res}° \n - 耗费时间：{elapse:.4f}s')
def get_cls_res(img_file_buffer):
image = Image.open(img_file_buffer)
img = np.array(image)
img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
orientation_res, elapse = orientation_engine(img_array)
return image, orientation_res, elapse
def layout_part():
st.markdown('##### 文档图像版面分析')
img = image_select(
label="示例图（点击选择）",
images=[
'images/layout/layout.jpg',
'images/layout/layout1.jpg',
'images/layout/layout2.jpg',
],
key='layout_select'
)
img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
key='layout',
label_visibility='collapsed')
layout_col1, layout_col2 = st.columns([5, 5])
img_empty = layout_col1.empty()
image, drawed_img = get_layout_res(img)
if img_file_buffer:
image, drawed_img = get_layout_res(img_file_buffer)
img_empty.image(image, use_column_width=True)
layout_col2.image(drawed_img, use_column_width=True)
def get_layout_res(img_file_buffer):
image = Image.open(img_file_buffer)
img = np.array(image)
img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
layout_res, _ = layout_engine(img_array)
drawed_img = vis_layout(img, layout_res)
return image, drawed_img
def vis_layout(img: np.ndarray, layout_res: list) -> None:
tmp_img = copy.deepcopy(img)
for v in layout_res:
bbox = np.round(v['bbox']).astype(np.int32)
label = v['label']
start_point = (bbox[0], bbox[1])
end_point = (bbox[2], bbox[3])
cv2.rectangle(tmp_img, start_point, end_point, (0, 0, 0), 2)
cv2.putText(tmp_img, label, start_point,
cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2)
return tmp_img
def table_part():
st.markdown('##### 表格还原')
img = image_select(
label="示例图（点击选择）",
images=[
'images/table/table.jpg',
'images/table/table1.png',
'images/table/table2.png',
],
key='table_select'
)
img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
key='table',
label_visibility='collapsed')
table_col1, table_col2 = st.columns([5, 5])
img_empty = table_col1.empty()
image, table_html_str = get_table_res(img)
if img_file_buffer:
image, table_html_str = get_table_res(img_file_buffer)
img_empty.image(image, use_column_width=True)
table_col2.markdown(table_html_str, unsafe_allow_html=True)
def get_table_res(img_file_buffer):
image = Image.open(img_file_buffer)
img = np.array(image)
img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
table_html_str, _, _ = table_engine(img_array)
table_html_str = vis_table(table_html_str)
return image, table_html_str
def vis_table(table_res) -> str:
style_res = ''''''
prefix_table, suffix_table = table_res.split('')
new_table_res = f'{prefix_table}{style_res}{suffix_table}'
return new_table_res
if __name__ == '__main__':
st.markdown("", unsafe_allow_html=True)
st.markdown("""
""", unsafe_allow_html=True)
img_suffix = ["png", "jpg", "jpeg"]
orientation_part()
layout_part()
table_part()