add adapter for OpenCV

This commit is contained in:
Yi Ge 2019-12-15 21:38:17 +08:00
parent 720c9d479f
commit 9360ebdd40
2 changed files with 25 additions and 13 deletions

15
videocr/opencv_adapter.py Normal file
View File

@ -0,0 +1,15 @@
import cv2
class Capture:
def __init__(self, video_path):
self.path = video_path
def __enter__(self):
self.cap = cv2.VideoCapture(self.path)
if not self.cap.isOpened():
raise IOError('Can not open video {}.'.format(self.path))
return self.cap
def __exit__(self, exc_type, exc_value, traceback):
self.cap.release()

View File

@ -6,6 +6,7 @@ import cv2
from . import constants from . import constants
from .models import PredictedFrame, PredictedSubtitle from .models import PredictedFrame, PredictedSubtitle
from .opencv_adapter import Capture
class Video: class Video:
@ -20,13 +21,10 @@ class Video:
def __init__(self, path: str): def __init__(self, path: str):
self.path = path self.path = path
v = cv2.VideoCapture(path) with Capture(path) as v:
if not v.isOpened(): self.num_frames = int(v.get(cv2.CAP_PROP_FRAME_COUNT))
raise IOError('can not open video format {}'.format(path)) self.fps = v.get(cv2.CAP_PROP_FPS)
self.num_frames = int(v.get(cv2.CAP_PROP_FRAME_COUNT)) self.height = int(v.get(cv2.CAP_PROP_FRAME_HEIGHT))
self.fps = v.get(cv2.CAP_PROP_FPS)
self.height = int(v.get(cv2.CAP_PROP_FRAME_HEIGHT))
v.release()
def run_ocr(self, lang: str, time_start: str, time_end: str, def run_ocr(self, lang: str, time_start: str, time_end: str,
conf_threshold:int, use_fullframe: bool) -> None: conf_threshold:int, use_fullframe: bool) -> None:
@ -41,13 +39,12 @@ class Video:
num_ocr_frames = ocr_end - ocr_start num_ocr_frames = ocr_end - ocr_start
# get frames from ocr_start to ocr_end # get frames from ocr_start to ocr_end
v = cv2.VideoCapture(self.path) with Capture(self.path) as v, multiprocessing.Pool() as pool:
v.set(cv2.CAP_PROP_POS_FRAMES, ocr_start) v.set(cv2.CAP_PROP_POS_FRAMES, ocr_start)
frames = (v.read()[1] for _ in range(num_ocr_frames)) frames = (v.read()[1] for _ in range(num_ocr_frames))
# perform ocr to frames in parallel # perform ocr to frames in parallel
with Pool() as pool: it_ocr = pool.imap(self._image_to_data, frames, chunksize=10)
it_ocr = pool.imap(self._single_frame_ocr, frames, chunksize=10)
self.pred_frames = [ self.pred_frames = [
PredictedFrame(i + ocr_start, data, conf_threshold) PredictedFrame(i + ocr_start, data, conf_threshold)
for i, data in enumerate(it_ocr)] for i, data in enumerate(it_ocr)]