diff --git a/videocr/api.py b/videocr/api.py index b182d43..a297f41 100644 --- a/videocr/api.py +++ b/videocr/api.py @@ -8,16 +8,17 @@ from .video import Video def get_subtitles( video_path: str, lang='eng', time_start='0:00', time_end='', conf_threshold=65, sim_threshold=90, use_fullframe=False) -> str: - # download tesseract data file to ~/tessdata if necessary - fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(lang) - if not fpath.is_file(): - if lang == 'eng': - url = constants.ENG_URL - else: - url = constants.TESSDATA_URL.format(lang) - constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True) - with urlopen(url) as res, open(fpath, 'w+b') as f: - shutil.copyfileobj(res, f) + # download tesseract data files to ~/tessdata if necessary + 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.run_ocr(lang, time_start, time_end, conf_threshold, use_fullframe) @@ -28,7 +29,7 @@ def save_subtitles_to_file( video_path: str, file_path='subtitle.srt', lang='eng', time_start='0:00', time_end='', conf_threshold=65, sim_threshold=90, use_fullframe=False) -> None: - with open(file_path, 'w+') as f: + with open(file_path, 'w+', encoding='utf-8') as f: f.write(get_subtitles( video_path, lang, time_start, time_end, conf_threshold, sim_threshold, use_fullframe)) diff --git a/videocr/constants.py b/videocr/constants.py index cfe0927..f13eb24 100644 --- a/videocr/constants.py +++ b/videocr/constants.py @@ -2,6 +2,6 @@ import pathlib TESSDATA_DIR = pathlib.Path.home() / 'tessdata' -ENG_URL = 'https://github.com/tesseract-ocr/tessdata_fast/raw/master/eng.traineddata' +TESSDATA_URL = 'https://github.com/tesseract-ocr/tessdata_fast/raw/master/{}.traineddata' -TESSDATA_URL = 'https://github.com/tesseract-ocr/tessdata_best/raw/master/script/{}.traineddata' +TESSDATA_SCRIPT_URL = 'https://github.com/tesseract-ocr/tessdata_best/raw/master/script/{}.traineddata'