# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
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(
"",
unsafe_allow_html=True,
)
st.markdown(
"""
""",
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")
save_file_dir = None
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_file_dir.resolve().parent))