app.py

# -*- 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("<h1 style='text-align: center;'><a href='https://github.com/RapidAI/PaddleOCRModelConverter' style='text-decoration: none'>PaddleOCRModelConverter</a></h1>", unsafe_allow_html=True)
st.markdown("""
<p>
    <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://pypi.org/project/paddleocr_convert/"><img alt="PyPI" src="https://img.shields.io/pypi/v/paddleocr_convert"></a>
    <a href="https://pepy.tech/project/paddleocr_convert"><img src="https://static.pepy.tech/personalized-badge/paddleocr_convert?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
</p>

""", 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))