SWHL/RapidVideOCR

1# -*- encoding: utf-8 -*-2# @Author: SWHL3# @Contact: liekkaskono@163.com4import shutil5import zipfile6from pathlib import Path7 8import streamlit as st9from rapid_videocr import RapidVideOCR10 11 12def mkdir(dir_path):13    Path(dir_path).mkdir(parents=True, exist_ok=True)14 15 16st.markdown(17    "<h1 style='text-align: center;'><a href='https://github.com/SWHL/RapidVideOCR' style='text-decoration: none'>RapidVideOCR</a></h1>",18    unsafe_allow_html=True,19)20st.markdown(21    """22<p align="left">23    <a href="https://colab.research.google.com/github/SWHL/RapidVideOCR/blob/75dae6e9804dec6e61bef98334601908dc9ec9fb/assets/RapidVideOCRDemo.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg"></a>24    <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>25    <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>26    <a href="https://pypi.org/project/rapid-videocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid_videocr"></a>27    <a href="https://github.com/SWHL/RapidVideOCR/stargazers"><img src="https://img.shields.io/github/stars/SWHL/RapidVideOCR?color=ccf"></a>28    <a href="https://pepy.tech/project/rapid-videocr"><img src="https://static.pepy.tech/personalized-badge/rapid-videocr?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>29    <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>30    <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>31</p>32""",33    unsafe_allow_html=True,34)35 36img_suffix = [".png", ".jpeg", ".jpg"]37 38st.markdown("#### Whether to use the concat rec:")39is_concat_res = st.checkbox("is_concat_rec")40if is_concat_res:41    number = st.number_input(42        "Insert concat batch num",43        value=10,44        placeholder="Default is 10",45        min_value=1,46        max_value=20,47    )48 49extractor = RapidVideOCR(50    is_concat_rec=is_concat_res, out_format='srt', is_print_console=False51)52 53save_dir = Path("images")54mkdir(save_dir)55output_dir = Path("outputs")56 57st.markdown("#### Upload ZIP files of RGBImages or TXTImages from VideoSubFinder:")58with st.form('my-form', clear_on_submit=True):59    uploaded_file = st.file_uploader(60        "Upload a ZIP file containg RGBImages or TXTImages folder.",61        type=["zip"],62        label_visibility="collapsed",63    )64    submitted = st.form_submit_button('Upload')65 66    if submitted and uploaded_file is not None:67        with zipfile.ZipFile(uploaded_file, "r") as zip_ref:68            file_list = zip_ref.namelist()69            for file_path in file_list:70                if Path(file_path).suffix not in img_suffix:71                    continue72 73                # 写到本地临时目录下74                save_full_path = save_dir / file_path75 76                save_file_dir = save_full_path.parent77                save_file_dir.mkdir(parents=True, exist_ok=True)78                with open(save_full_path, "wb") as f:79                    f.write(zip_ref.read(file_path))80        with st.spinner('Extracting...'):81            extractor(save_file_dir, output_dir)82 83srt_path = output_dir / "result.srt"84if srt_path is not None and srt_path.exists():85    st.toast('Success!', icon='🎉')86    with open(srt_path, "rb") as f:87        is_download_srt = st.download_button(88            "Download SRT File", data=f, file_name=Path(srt_path).name89        )90        if is_download_srt:91            shutil.rmtree(str(srt_path.resolve().parent))92            shutil.rmtree(str(save_dir.resolve()))93