# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import shutil
from pathlib import Path
import streamlit as st
from paddleocr_convert import PaddleOCRModelConvert
converter = PaddleOCRModelConvert()
st.markdown("", unsafe_allow_html=True)
st.markdown("""
""", unsafe_allow_html=True)
st.markdown('戳这里查看支持模型列表：[PaddleOCR Models](https://github.com/PaddlePaddle/PaddleOCR/blob/40c56628fda416e1c8710eb19e4b260536902520/doc/doc_ch/models_list.md)')
st.markdown('⚠️注意：不支持slim量化后的模型')
default_url = 'e.g. https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar'
url = st.text_input('输入模型地址：', help=default_url)
default_txt_path = 'e.g. https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.6/ppocr/utils/ppocr_keys_v1.txt'
txt_url = st.text_input('输入txt地址（文本识别模型时，必选）：', help=default_txt_path)
save_dir = 'models'
save_onnx_path = ''
is_convert = st.button('Convert')
if is_convert:
if not url or len(url) <= 0:
st.error('模型链接不能为空！')
st.stop()
if 'rec' in url and not txt_url:
st.error('识别模型对应字典不能为空')
st.stop()
with st.spinner("正在转换，请稍等........"):
save_onnx_path = converter(url, save_dir,
txt_path=txt_url,
is_del_raw=True)
st.success('转换成功，点击Download下载!', icon="✅")
if save_onnx_path:
with open(save_onnx_path, 'rb') as file:
is_download = st.download_button('Download', data=file,
file_name=Path(save_onnx_path).name)
if is_download:
shutil.rmtree(str(Path(save_onnx_path).resolve().parent))