From a5e6845a1bc3d3497bda2e87b87bbeaf2541b3f2 Mon Sep 17 00:00:00 2001
From: Yi Ge <me@yige.ch>
Date: Mon, 29 Apr 2019 03:04:06 +0200
Subject: [PATCH] move tessdata dir to ~/tessdata

---
 videocr/api.py       | 8 +++-----
 videocr/constants.py | 4 ++++
 videocr/video.py     | 5 +++--
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/videocr/api.py b/videocr/api.py
index 3ef693d..c581890 100644
--- a/videocr/api.py
+++ b/videocr/api.py
@@ -1,6 +1,4 @@
 from urllib.request import urlopen
-import pathlib
-import os
 import shutil
 
 from . import constants
@@ -9,14 +7,14 @@ 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))
+    # 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)
-        os.makedirs('tessdata', exist_ok=True)
+        constants.TESSDATA_DIR.mkdir(parents=True, exist_ok=True)
         with urlopen(url) as res, open(fpath, 'w+b') as f:
             shutil.copyfileobj(res, f)
 
diff --git a/videocr/constants.py b/videocr/constants.py
index 91d8453..cfe0927 100644
--- a/videocr/constants.py
+++ b/videocr/constants.py
@@ -1,3 +1,7 @@
+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_best/raw/master/script/{}.traineddata'
diff --git a/videocr/video.py b/videocr/video.py
index de9c000..fb2fe04 100644
--- a/videocr/video.py
+++ b/videocr/video.py
@@ -4,6 +4,7 @@ import datetime
 import pytesseract
 import cv2
 
+from . import constants
 from .models import PredictedFrame, PredictedSubtitle
 
 
@@ -70,8 +71,8 @@ class Video:
     def _single_frame_ocr(self, img) -> str:
         if not self.use_fullframe:
             # only use bottom half of the frame by default
-            img = img[img.shape[0] // 2:, :]
-        config = r'--tessdata-dir "tessdata"'
+            img = img[self.height // 2:, :]
+        config = '--tessdata-dir "{}"'.format(constants.TESSDATA_DIR)
         return pytesseract.image_to_data(img, lang=self.lang, config=config)
 
     def get_subtitles(self) -> str: