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()