all repos — videocr @ master

Extract hardcoded subtitles from videos using machine learning

videocr/utils.py (view raw)

 1from urllib.request import urlopen
 2import shutil
 3import datetime
 4
 5from . import constants
 6
 7
 8# download language data files to ~/tessdata if necessary
 9def download_lang_data(lang: str):
10    constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
11
12    for lang_name in lang.split('+'):
13        filepath = constants.TESSDATA_DIR / '{}.traineddata'.format(lang_name)
14        if not filepath.is_file():
15            # download needed file
16            if lang_name[0].isupper():
17                url = constants.TESSDATA_SCRIPT_URL.format(lang_name)
18            else:
19                url = constants.TESSDATA_URL.format(lang_name)
20
21            with urlopen(url) as res, open(filepath, 'w+b') as f:
22                shutil.copyfileobj(res, f)
23
24
25# convert time string to frame index
26def get_frame_index(time_str: str, fps: float):
27    t = time_str.split(':')
28    t = list(map(float, t))
29    if len(t) == 3:
30        td = datetime.timedelta(hours=t[0], minutes=t[1], seconds=t[2])
31    elif len(t) == 2:
32        td = datetime.timedelta(minutes=t[0], seconds=t[1])
33    else:
34        raise ValueError(
35            'Time data "{}" does not match format "%H:%M:%S"'.format(time_str))
36    index = int(td.total_seconds() * fps)
37    return index
38
39
40# convert frame index into SRT timestamp
41def get_srt_timestamp(frame_index: int, fps: float):
42    td = datetime.timedelta(seconds=frame_index / fps)
43    ms = td.microseconds // 1000
44    m, s = divmod(td.seconds, 60)
45    h, m = divmod(m, 60)
46    return '{:02d}:{:02d}:{:02d},{:03d}'.format(h, m, s, ms)