auto download tesseract data file

videocr/api.py

@@ -1,9 +1,25 @@
+from urllib.request import urlopen
+import pathlib
+import os
+import shutil
 
+from . import constants
 from .video import Video
 
 
 def get_subtitles(video_path: str, lang='eng',
                   time_start='0:00', time_end='', use_fullframe=False) -> str:
+    # download tesseract data file to ./tessdata/ if necessary
+    fpath = pathlib.Path('tessdata/{}.traineddata'.format(lang))
+    if not fpath.is_file():
+        if lang == 'eng':
+            url = constants.ENG_URL
+        else:
+            url = constants.TESSDATA_URL.format(lang)
+        os.makedirs('tessdata', exist_ok=True)
+        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, use_fullframe)
     return v.get_subtitles()

videocr/constants.py

@@ -0,0 +1,3 @@
+ENG_URL = 'https://github.com/tesseract-ocr/tessdata_fast/raw/master/eng.traineddata'
+
+TESSDATA_URL = 'https://github.com/tesseract-ocr/tessdata_best/raw/master/script/{}.traineddata'

videocr/models.py

@@ -4,10 +4,6 @@ from dataclasses import dataclass
 from fuzzywuzzy import fuzz
 
 
-CONF_THRESHOLD = 60
-# word predictions with lower confidence will be filtered out
-
-
 @dataclass
 class PredictedWord:
     __slots__ = 'confidence', 'text'
@@ -21,7 +17,7 @@ class PredictedFrame:
     confidence: int  # total confidence of all words
     text: str
 
-    def __init__(self, index, pred_data: str):
+    def __init__(self, index: int, pred_data: str, conf_threshold=70):
         self.index = index
         self.words = []
 
@@ -41,7 +37,8 @@ class PredictedFrame:
                 if self.words and self.words[-1].text != '\n':
                     self.words.append(PredictedWord(0, '\n'))
 
-            if conf >= CONF_THRESHOLD:
+            # word predictions with low confidence will be filtered out
+            if conf >= conf_threshold:
                 self.words.append(PredictedWord(conf, text))
 
         self.confidence = sum(word.confidence for word in self.words)

videocr/video.py

@@ -71,7 +71,8 @@ class Video:
         if not self.use_fullframe:
             # only use bottom half of the frame by default
             img = img[img.shape[0] // 2:, :]
-        return pytesseract.image_to_data(img, lang=self.lang)
+        config = r'--tessdata-dir "tessdata"'
+        return pytesseract.image_to_data(img, lang=self.lang, config=config)
 
     def get_subtitles(self) -> str:
         self._generate_subtitles()