Spaces:
Running
Running
File size: 3,627 Bytes
7561011 9e7a178 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# -*- 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()))
|