auto download tesseract data file

This commit is contained in:
Yi Ge 2019-04-28 17:33:16 +02:00
parent bccdcc02fc
commit fba35f0108
4 changed files with 24 additions and 7 deletions

View File

@ -1,9 +1,25 @@
from urllib.request import urlopen
import pathlib
import os
import shutil
from . import constants
from .video import Video 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
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 = Video(video_path)
v.run_ocr(lang, time_start, time_end, use_fullframe) v.run_ocr(lang, time_start, time_end, use_fullframe)
return v.get_subtitles() return v.get_subtitles()

3
videocr/constants.py Normal file
View File

@ -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'

View File

@ -4,10 +4,6 @@ from dataclasses import dataclass
from fuzzywuzzy import fuzz from fuzzywuzzy import fuzz
CONF_THRESHOLD = 60
# word predictions with lower confidence will be filtered out
@dataclass @dataclass
class PredictedWord: class PredictedWord:
__slots__ = 'confidence', 'text' __slots__ = 'confidence', 'text'
@ -21,7 +17,7 @@ class PredictedFrame:
confidence: int # total confidence of all words confidence: int # total confidence of all words
text: str text: str
def __init__(self, index, pred_data: str): def __init__(self, index: int, pred_data: str, conf_threshold=70):
self.index = index self.index = index
self.words = [] self.words = []
@ -41,7 +37,8 @@ class PredictedFrame:
if self.words and self.words[-1].text != '\n': if self.words and self.words[-1].text != '\n':
self.words.append(PredictedWord(0, '\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.words.append(PredictedWord(conf, text))
self.confidence = sum(word.confidence for word in self.words) self.confidence = sum(word.confidence for word in self.words)

View File

@ -71,7 +71,8 @@ class Video:
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[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: def get_subtitles(self) -> str:
self._generate_subtitles() self._generate_subtitles()