Spaces:
Running
Running
| # -*- 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( | |
| "<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())) | |