sdk/platform-tools/systrace/catapult/telemetry/telemetry/internal/image_processing/video_file_frame_generator.py

   1 # Copyright 2014 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 from telemetry.internal.image_processing import frame_generator
   6 from telemetry.internal.util import external_modules
   7
   8 cv2 = external_modules.ImportRequiredModule('cv2')
   9
  10
  11 class VideoFileFrameGenerator(frame_generator.FrameGenerator):
  12   """Provides a Frame Generator for a video file.
  13
  14   Sample Usage:
  15     generator = VideoFileFrameGenerator(sys.argv[1]).GetGenerator()
  16     for frame in generator:
  17       # Do something
  18
  19   Attributes:
  20     _capture: The openCV video capture.
  21     _frame_count: The number of frames in the video capture.
  22     _frame_index: The frame number of the current frame.
  23     _timestamp: The timestamp of the current frame.
  24     _dimensions: The dimensions of the video capture."""
  25   def __init__(self, video_filename, start_frame_index=0):
  26     """Initializes the VideoFileFrameGenerator object.
  27
  28     Args:
  29       video_filename: str, The path to the video file.
  30       start_frame_index: int, The number of frames to skip at the start of the
  31           file.
  32
  33     Raises:
  34       FrameReadError: A read error occurred during initialization."""
  35     self._capture = cv2.VideoCapture(video_filename)
  36     self._frame_count = int(self._capture.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
  37     self._frame_index = -1
  38     self._timestamp = 0
  39     width = self._capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)
  40     height = self._capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)
  41     self._dimensions = (int(width), int(height))
  42     if self._frame_count <= start_frame_index:
  43       raise frame_generator.FrameReadError('Not enough frames in capture.')
  44     while self._frame_index < start_frame_index - 1:
  45       self._ReadFrame(True)
  46
  47     super(self.__class__, self).__init__()
  48
  49   def _ReadFrame(self, skip_decode=False):
  50     """Reads the next frame, updates attributes.
  51
  52     Args:
  53       skip_decode: Whether or not to skip decoding. Useful for seeking.
  54
  55     Returns:
  56       The frame if not EOF, 'None' if EOF.
  57
  58     Raises:
  59       FrameReadError: Unexpectedly failed to read a frame from the capture."""
  60     if self._frame_index >= self._frame_count - 1:
  61       return None
  62     self._timestamp = self._capture.get(cv2.cv.CV_CAP_PROP_POS_MSEC)
  63     if skip_decode:
  64       ret = self._capture.grab()
  65       frame = None
  66     else:
  67       ret, frame = self._capture.read()
  68     if not ret:
  69       raise frame_generator.FrameReadError(
  70           'Failed to read frame from capture.')
  71     self._frame_index += 1
  72     return frame
  73
  74   # OVERRIDE
  75   def _CreateGenerator(self):
  76     while True:
  77       frame = self._ReadFrame()
  78       if frame is None:
  79         break
  80       yield frame
  81
  82   # OVERRIDE
  83   @property
  84   def CurrentTimestamp(self):
  85     return self._timestamp
  86
  87   # OVERRIDE
  88   @property
  89   def CurrentFrameNumber(self):
  90     return self._frame_index
  91
  92   # OVERRIDE
  93   @property
  94   def Dimensions(self):
  95     return self._dimensions