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
5from telemetry.internal.image_processing import frame_generator
6from telemetry.internal.util import external_modules
7
8cv2 = external_modules.ImportRequiredModule('cv2')
9
10
11class 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
96