From fba35f010840753d71b2bafc9447f5607031eb7f Mon Sep 17 00:00:00 2001
From: Yi Ge <me@yige.ch>
Date: Sun, 28 Apr 2019 17:33:16 +0200
Subject: [PATCH] auto download tesseract data file

---
 videocr/api.py       | 16 ++++++++++++++++
 videocr/constants.py |  3 +++
 videocr/models.py    |  9 +++------
 videocr/video.py     |  3 ++-
 4 files changed, 24 insertions(+), 7 deletions(-)
 create mode 100644 videocr/constants.py

diff --git a/videocr/api.py b/videocr/api.py
index f253638..3ef693d 100644
--- a/videocr/api.py
+++ b/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()
diff --git a/videocr/constants.py b/videocr/constants.py
new file mode 100644
index 0000000..91d8453
--- /dev/null
+++ b/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'
diff --git a/videocr/models.py b/videocr/models.py
index f787712..5b5549f 100644
--- a/videocr/models.py
+++ b/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)
diff --git a/videocr/video.py b/videocr/video.py
index 2485e5d..de9c000 100644
--- a/videocr/video.py
+++ b/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()