1from urllib.request import urlopen
 2import shutil
 3
 4from . import constants
 5from .video import Video
 6
 7
 8def get_subtitles(
 9        video_path: str, lang='eng', time_start='0:00', time_end='',
10        conf_threshold=65, sim_threshold=90, use_fullframe=False) -> str:
11    # download tesseract data files to ~/tessdata if necessary
12    constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
13    for fname in lang.split('+'):
14        fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(fname)
15        if not fpath.is_file():
16            if fname[0].isupper():
17                url = constants.TESSDATA_SCRIPT_URL.format(fname)
18            else:
19                url = constants.TESSDATA_URL.format(fname)
20            with urlopen(url) as res, open(fpath, 'w+b') as f:
21                shutil.copyfileobj(res, f)
22
23    v = Video(video_path)
24    v.run_ocr(lang, time_start, time_end, conf_threshold, use_fullframe)
25    return v.get_subtitles(sim_threshold)
26
27
28def save_subtitles_to_file(
29        video_path: str, file_path='subtitle.srt', lang='eng',
30        time_start='0:00', time_end='', conf_threshold=65, sim_threshold=90,
31        use_fullframe=False) -> None:
32    with open(file_path, 'w+', encoding='utf-8') as f:
33        f.write(get_subtitles(
34            video_path, lang, time_start, time_end, conf_threshold,
35            sim_threshold, use_fullframe))