M videocr/api.py → videocr/api.py

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

M videocr/constants.py → videocr/constants.py

@@ -2,6 +2,6 @@ import pathlib
 
 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'