app.py

# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: [email protected]
import shutil
import zipfile
from pathlib import Path

import streamlit as st
from rapid_videocr import RapidVideOCR


def mkdir(dir_path):
    Path(dir_path).mkdir(parents=True, exist_ok=True)


st.markdown(
    "<h1 style='text-align: center;'><a href='https://github.com/SWHL/RapidVideOCR' style='text-decoration: none'>RapidVideOCR</a></h1>",
    unsafe_allow_html=True,
)
st.markdown(
    """
<p align="left">
    <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>
    <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-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/rapid-videocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid_videocr"></a>
    <a href="https://github.com/SWHL/RapidVideOCR/stargazers"><img src="https://img.shields.io/github/stars/SWHL/RapidVideOCR?color=ccf"></a>
    <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>
    <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
    <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
</p>
""",
    unsafe_allow_html=True,
)

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

st.markdown("#### Whether to use the concat rec:")
is_concat_res = st.checkbox("is_concat_rec")
if is_concat_res:
    number = st.number_input(
        "Insert concat batch num",
        value=10,
        placeholder="Default is 10",
        min_value=1,
        max_value=20,
    )

extractor = RapidVideOCR(
    is_concat_rec=is_concat_res, out_format='srt', is_print_console=False
)

save_dir = Path("images")
mkdir(save_dir)
output_dir = Path("outputs")

st.markdown("#### Upload ZIP files of RGBImages or TXTImages from VideoSubFinder:")
with st.form('my-form', clear_on_submit=True):
    uploaded_file = st.file_uploader(
        "Upload a ZIP file containg RGBImages or TXTImages folder.",
        type=["zip"],
        label_visibility="collapsed",
    )
    submitted = st.form_submit_button('Upload')

    if submitted and uploaded_file is not None:
        with zipfile.ZipFile(uploaded_file, "r") as zip_ref:
            file_list = zip_ref.namelist()
            for file_path in file_list:
                if Path(file_path).suffix not in img_suffix:
                    continue

                # 写到本地临时目录下
                save_full_path = save_dir / file_path

                save_file_dir = save_full_path.parent
                save_file_dir.mkdir(parents=True, exist_ok=True)
                with open(save_full_path, "wb") as f:
                    f.write(zip_ref.read(file_path))
        with st.spinner('Extracting...'):
            extractor(save_file_dir, output_dir)

srt_path = output_dir / "result.srt"
if srt_path is not None and srt_path.exists():
    st.toast('Success!', icon='🎉')
    with open(srt_path, "rb") as f:
        is_download_srt = st.download_button(
            "Download SRT File", data=f, file_name=Path(srt_path).name
        )
        if is_download_srt:
            shutil.rmtree(str(srt_path.resolve().parent))
            shutil.rmtree(str(save_dir.resolve()))