1from urllib.request import urlopen
 2import pathlib
 3import os
 4import shutil
 5
 6from . import constants
 7from .video import Video
 8
 9
10def get_subtitles(video_path: str, lang='eng',
11                  time_start='0:00', time_end='', use_fullframe=False) -> str:
12    # download tesseract data file to ./tessdata/ if necessary
13    fpath = pathlib.Path('tessdata/{}.traineddata'.format(lang))
14    if not fpath.is_file():
15        if lang == 'eng':
16            url = constants.ENG_URL
17        else:
18            url = constants.TESSDATA_URL.format(lang)
19        os.makedirs('tessdata', exist_ok=True)
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, use_fullframe)
25    return v.get_subtitles()
26
27
28def save_subtitles_to_file(
29        video_path: str, file_path='subtitle.srt', lang='eng',
30        time_start='0:00', time_end='', use_fullframe=False) -> None:
31    with open(file_path, 'w+') as f:
32        f.write(get_subtitles(
33            video_path, lang, time_start, time_end, use_fullframe))