app.py

# -*- 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 = '''<style>td {border-left: 1px solid;border-bottom:1px solid;}
                   table, th {border-top:1px solid;font-size: 10px;
                   border-collapse: collapse;border-right: 1px solid;}
                </style>'''
    prefix_table, suffix_table = table_res.split('<body>')
    new_table_res = f'{prefix_table}{style_res}<body>{suffix_table}'
    return new_table_res


if __name__ == '__main__':
    st.markdown("<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidStructure' style='text-decoration: none'>Rapid Structure</a></h1>", unsafe_allow_html=True)
    st.markdown("""
    <p align="left">
        <a href=""><img src="https://img.shields.io/badge/Python->=3.7,<=3.10-aff.svg"></a>
        <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></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=rapid-layout"></a>
        <a href="https://pepy.tech/project/rapid-orientation"><img src="https://static.pepy.tech/personalized-badge/rapid-orientation?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=rapid-orientation"></a>
        <a href="https://pepy.tech/project/rapid-table"><img src="https://static.pepy.tech/personalized-badge/rapid-table?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=rapid-table"></a>
    </p>
    """, unsafe_allow_html=True)

    img_suffix = ["png", "jpg", "jpeg"]

    orientation_part()
    layout_part()
    table_part()