calculate PredictedSubtitle.text early
Yi Ge me@yige.ch
Fri, 26 Apr 2019 00:32:47 +0200
2 files changed,
18 insertions(+),
16 deletions(-)
M
videocr/models.py
→
videocr/models.py
@@ -62,12 +62,12 @@
def __init__(self, frames: List[PredictedFrame]): self.frames = [f for f in frames if f.confidence > 0] - @property - def text(self) -> str: if self.frames: conf_max = max(f.confidence for f in self.frames) - return next(f.text for f in self.frames if f.confidence == conf_max) - return '' + self.text = next(f.text for f in self.frames + if f.confidence == conf_max) + else: + self.text = '' @property def index_start(self) -> int:
M
videocr/video.py
→
videocr/video.py
@@ -46,12 +46,22 @@
v.release() def get_subtitles(self) -> str: + self._generate_subtitles() + return ''.join( + '{}\n{} --> {}\n{}\n'.format( + i, + self._srt_timestamp(sub.index_start), + self._srt_timestamp(sub.index_end), + sub.text) + for i, sub in enumerate(self.pred_subs)) + + def _generate_subtitles(self) -> None: + self.pred_subs = [] + if self.pred_frames is None: raise AttributeError( 'Please call self.run_ocr() first to generate ocr of frames') - self.pred_subs = [] - # divide ocr of frames into subtitle paragraphs using sliding window WIN_BOUND = int(self.fps / 2) # 1/2 sec sliding window boundary bound = WIN_BOUND@@ -75,18 +85,10 @@ bound = WIN_BOUND
j += 1 + # also handle the last remaining frames if i < self.num_frames - 1: self._append_sub(PredictedSubtitle(self.pred_frames[i:])) - for i, sub in enumerate(self.pred_subs): - print('{}\n{} --> {}\n{}\n'.format( - i, - self._srt_timestamp(sub.index_start), - self._srt_timestamp(sub.index_end), - sub.text)) - - return '' - def _append_sub(self, sub: PredictedSubtitle) -> None: if len(sub.text) == 0: return@@ -101,7 +103,7 @@ self.pred_subs.append(sub)
def _srt_timestamp(self, frame_index) -> str: time = str(datetime.timedelta(seconds=frame_index / self.fps)) - return time.replace('.', ',') # srt uses comma as fractional separator + return time.replace('.', ',') # srt uses comma, not dot time_start = timeit.default_timer()