move tessdata dir to ~/tessdata
This commit is contained in:
parent
fba35f0108
commit
a5e6845a1b
|
@ -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)
|
||||||
|
|
||||||
|
|
|
@ -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'
|
||||||
|
|
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue