# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import cv2
import numpy as np
import streamlit as st
from PIL import Image
from rapid_orientation import RapidOrientation

orientation_engine = RapidOrientation()


if __name__ == '__main__':
    st.markdown("<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Orientation.md' style='text-decoration: none'>rapid-orientation</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-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>
    </p>

    #### 简介
    - 该部分主要是做含文字图像方向分类模型。模型来源：[PaddleClas 含文字图像方向分类模型](https://github.com/PaddlePaddle/PaddleClas/blob/177e4be74639c0960efeae2c5166d3226c9a02eb/docs/zh_CN/models/PULC/PULC_text_image_orientation.md)

        | 模型类型  |        模型名称         | 模型大小 |                           支持种类                           |
        |:---:|:---:|:---:|:---:|
        |   四方向分类   |   `rapid_orientation.onnx`   |  6.5M | `0 90 180 270`|

    #### 试一下
    """, unsafe_allow_html=True)

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

    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix)
    col1, col2 = st.columns([5, 5])

    img_empty = col1.empty()
    if img_file_buffer:
        image = Image.open(img_file_buffer)
        img = np.array(image)
        img_empty.image(image, use_column_width=True)

        img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        orientation_res, elapse = orientation_engine(img)

        col2.markdown(f'- 方向分类结果：{orientation_res}° \n - 耗费时间：{elapse:.4f}s')