move download_lang_data to utils.py

This commit is contained in:
Yi Ge 2019-12-15 20:56:09 +08:00
parent 8a56cbf746
commit 720c9d479f
2 changed files with 23 additions and 15 deletions

View File

@ -1,24 +1,11 @@
from urllib.request import urlopen from . import utils
import shutil
from . import constants
from .video import Video from .video import Video
def get_subtitles( def get_subtitles(
video_path: str, lang='eng', time_start='0:00', time_end='', video_path: str, lang='eng', time_start='0:00', time_end='',
conf_threshold=65, sim_threshold=90, use_fullframe=False) -> str: conf_threshold=65, sim_threshold=90, use_fullframe=False) -> str:
# download tesseract data files to ~/tessdata if necessary utils.download_lang_data(lang)
constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
for fname in lang.split('+'):
fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(fname)
if not fpath.is_file():
if fname[0].isupper():
url = constants.TESSDATA_SCRIPT_URL.format(fname)
else:
url = constants.TESSDATA_URL.format(fname)
with urlopen(url) as res, open(fpath, 'w+b') as f:
shutil.copyfileobj(res, f)
v = Video(video_path) v = Video(video_path)
v.run_ocr(lang, time_start, time_end, conf_threshold, use_fullframe) v.run_ocr(lang, time_start, time_end, conf_threshold, use_fullframe)

21
videocr/utils.py Normal file
View File

@ -0,0 +1,21 @@
from urllib.request import urlopen
import shutil
from . import constants
# download language data files to ~/tessdata if necessary
def download_lang_data(lang: str):
constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
for lang_name in lang.split('+'):
filepath = constants.TESSDATA_DIR / '{}.traineddata'.format(lang_name)
if not filepath.is_file():
# download needed file
if lang_name[0].isupper():
url = constants.TESSDATA_SCRIPT_URL.format(lang_name)
else:
url = constants.TESSDATA_URL.format(lang_name)
with urlopen(url) as res, open(filepath, 'w+b') as f:
shutil.copyfileobj(res, f)