move tessdata dir to ~/tessdata

This commit is contained in:
Yi Ge 2019-04-29 03:04:06 +02:00
parent fba35f0108
commit a5e6845a1b
3 changed files with 10 additions and 7 deletions

View File

@ -1,6 +1,4 @@
from urllib.request import urlopen from urllib.request import urlopen
import pathlib
import os
import shutil import shutil
from . import constants from . import constants
@ -9,14 +7,14 @@ from .video import Video
def get_subtitles(video_path: str, lang='eng', def get_subtitles(video_path: str, lang='eng',
time_start='0:00', time_end='', use_fullframe=False) -> str: time_start='0:00', time_end='', use_fullframe=False) -> str:
# download tesseract data file to ./tessdata/ if necessary # download tesseract data file to ~/tessdata if necessary
fpath = pathlib.Path('tessdata/{}.traineddata'.format(lang)) fpath = constants.TESSDATA_DIR / '{}.traineddata'.format(lang)
if not fpath.is_file(): if not fpath.is_file():
if lang == 'eng': if lang == 'eng':
url = constants.ENG_URL url = constants.ENG_URL
else: else:
url = constants.TESSDATA_URL.format(lang) 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: with urlopen(url) as res, open(fpath, 'w+b') as f:
shutil.copyfileobj(res, f) shutil.copyfileobj(res, f)

View File

@ -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' 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' TESSDATA_URL = 'https://github.com/tesseract-ocr/tessdata_best/raw/master/script/{}.traineddata'

View File

@ -4,6 +4,7 @@ import datetime
import pytesseract import pytesseract
import cv2 import cv2
from . import constants
from .models import PredictedFrame, PredictedSubtitle from .models import PredictedFrame, PredictedSubtitle
@ -70,8 +71,8 @@ class Video:
def _single_frame_ocr(self, img) -> str: def _single_frame_ocr(self, img) -> str:
if not self.use_fullframe: if not self.use_fullframe:
# only use bottom half of the frame by default # only use bottom half of the frame by default
img = img[img.shape[0] // 2:, :] img = img[self.height // 2:, :]
config = r'--tessdata-dir "tessdata"' config = '--tessdata-dir "{}"'.format(constants.TESSDATA_DIR)
return pytesseract.image_to_data(img, lang=self.lang, config=config) return pytesseract.image_to_data(img, lang=self.lang, config=config)
def get_subtitles(self) -> str: def get_subtitles(self) -> str: