Yun
/
videocr
forked from pradana.aumars/videocr
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

support combining multiple languages

This commit is contained in:
Yi Ge 2019-04-29 22:29:49 +02:00
parent 94e2ca5b8e
commit 04ad4597ff
2 changed files with 14 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
videocr/api.py
View File

@ -8,16 +8,17 @@ 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 file to ~/tessdata if necessary # download tesseract data files to ~/tessdata if necessary
fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(lang) constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
if not fpath.is_file(): for fname in lang.split('+'):
if lang == 'eng': fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(fname)
url = constants.ENG_URL if not fpath.is_file():
else: if fname[0].isupper():
url = constants.TESSDATA_URL.format(lang) url = constants.TESSDATA_SCRIPT_URL.format(fname)
constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True) else:
with urlopen(url) as res, open(fpath, 'w+b') as f: url = constants.TESSDATA_URL.format(fname)
shutil.copyfileobj(res, f) 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)
@ -28,7 +29,7 @@ def save_subtitles_to_file(
video_path: str, file_path='subtitle.srt', lang='eng', video_path: str, file_path='subtitle.srt', lang='eng',
time_start='0:00', time_end='', conf_threshold=65, sim_threshold=90, time_start='0:00', time_end='', conf_threshold=65, sim_threshold=90,
use_fullframe=False) -> None: 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( f.write(get_subtitles(
video_path, lang, time_start, time_end, conf_threshold, video_path, lang, time_start, time_end, conf_threshold,
sim_threshold, use_fullframe)) sim_threshold, use_fullframe))

4
videocr/constants.py
View File

@ -2,6 +2,6 @@ import pathlib
TESSDATA_DIR = pathlib.Path.home() / 'tessdata' 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'