support combining multiple languages
This commit is contained in:
parent
94e2ca5b8e
commit
04ad4597ff
@ -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))
|
||||||
|
@ -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'
|
||||||
|
Loading…
Reference in New Issue
Block a user