2019-04-24 21:18:31 +02:00
|
|
|
from __future__ import annotations
|
|
|
|
from concurrent import futures
|
|
|
|
import pytesseract
|
|
|
|
import cv2
|
|
|
|
import timeit
|
|
|
|
|
2019-04-25 01:40:46 +02:00
|
|
|
from .models import PredictedFrame, PredictedSubtitle
|
|
|
|
|
|
|
|
|
|
|
|
SUBTITLE_BOUND = 10
|
2019-04-24 21:18:31 +02:00
|
|
|
|
|
|
|
|
|
|
|
class Video:
|
|
|
|
path: str
|
|
|
|
lang: str
|
|
|
|
num_frames: int
|
2019-04-25 01:40:46 +02:00
|
|
|
pred_frames: List[PredictedFrame]
|
2019-04-24 21:18:31 +02:00
|
|
|
|
|
|
|
def __init__(self, path, lang):
|
|
|
|
self.path = path
|
|
|
|
self.lang = lang
|
|
|
|
v = cv2.VideoCapture(path)
|
|
|
|
self.num_frames = int(v.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
v.release()
|
|
|
|
|
2019-04-25 01:40:46 +02:00
|
|
|
def _single_frame_ocr(self, img) -> str:
|
|
|
|
img = img[img.shape[0] // 2:, :] # only use bottom half of the frame
|
2019-04-24 21:18:31 +02:00
|
|
|
data = pytesseract.image_to_data(img, lang=self.lang)
|
|
|
|
return data
|
|
|
|
|
2019-04-25 01:40:46 +02:00
|
|
|
def run_ocr(self) -> None:
|
2019-04-24 21:18:31 +02:00
|
|
|
v = cv2.VideoCapture(self.path)
|
2019-04-25 01:40:46 +02:00
|
|
|
frames = (v.read()[1] for _ in range(self.num_frames))
|
2019-04-24 21:18:31 +02:00
|
|
|
|
2019-04-25 01:40:46 +02:00
|
|
|
# perform ocr to all frames in parallel
|
2019-04-24 21:18:31 +02:00
|
|
|
with futures.ProcessPoolExecutor() as pool:
|
2019-04-25 01:40:46 +02:00
|
|
|
frames_ocr = pool.map(self._single_frame_ocr, frames, chunksize=10)
|
|
|
|
self.pred_frames = [PredictedFrame(i, data)
|
|
|
|
for i, data in enumerate(frames_ocr)]
|
2019-04-24 21:18:31 +02:00
|
|
|
|
|
|
|
v.release()
|
|
|
|
|
2019-04-25 01:40:46 +02:00
|
|
|
def get_subtitles(self) -> str:
|
|
|
|
if self.pred_frames is None:
|
|
|
|
raise AttributeError(
|
|
|
|
'Please call self.run_ocr() first to generate ocr of frames')
|
|
|
|
|
|
|
|
# divide ocr of frames into subtitle paragraphs using sliding window
|
|
|
|
i = 0
|
|
|
|
j = 1
|
|
|
|
bound = SUBTITLE_BOUND
|
|
|
|
while j < self.num_frames:
|
|
|
|
fi, fj = self.pred_frames[i], self.pred_frames[j]
|
|
|
|
|
|
|
|
if fi.is_similar_to(fj):
|
|
|
|
bound = SUBTITLE_BOUND
|
|
|
|
elif bound > 0:
|
|
|
|
bound -= 1
|
|
|
|
else:
|
|
|
|
# divide subtitle paragraphs
|
|
|
|
para_new = j - SUBTITLE_BOUND
|
|
|
|
print(PredictedSubtitle(self.pred_frames[i:para_new]).text)
|
|
|
|
i = para_new
|
|
|
|
j = i
|
|
|
|
bound = SUBTITLE_BOUND
|
|
|
|
|
|
|
|
j += 1
|
|
|
|
|
|
|
|
if i < self.num_frames - 1:
|
|
|
|
print(PredictedSubtitle(self.pred_frames[i:]).text)
|
|
|
|
|
|
|
|
return ''
|
|
|
|
|
2019-04-24 21:18:31 +02:00
|
|
|
|
|
|
|
time_start = timeit.default_timer()
|
|
|
|
v = Video('1.mp4', 'HanS')
|
|
|
|
v.run_ocr()
|
2019-04-25 01:40:46 +02:00
|
|
|
v.get_subtitles()
|
2019-04-24 21:18:31 +02:00
|
|
|
time_stop = timeit.default_timer()
|
|
|
|
print(time_stop - time_start)
|