[telemetry] Process streams of bitmaps via BitmapIter

tools/telemetry/telemetry/core/bitmapiter.py

Index: tools/telemetry/telemetry/core/bitmapiter.py
diff --git a/tools/telemetry/telemetry/core/bitmapiter.py b/tools/telemetry/telemetry/core/bitmapiter.py
new file mode 100644
index 0000000000000000000000000000000000000000..22444b7d7ff424dbb2e6c72e8c9fb0b6a0a028ef
--- /dev/null
+++ b/tools/telemetry/telemetry/core/bitmapiter.py
@@ -0,0 +1,175 @@
+# Copyright 2014 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import array
+import os
+import struct
+import subprocess
+
+from telemetry.core import bitmap
+from telemetry.core import util
+
+
+def _GetDimensions(video):
+  proc = subprocess.Popen(['avconv', '-i', video], stderr=subprocess.PIPE)
+  dimensions = None
+  output = ''
+  for line in proc.stderr.readlines():
+    output += line
+    if 'Video:' in line:
+      dimensions = line.split(',')[2]
+      dimensions = map(int, dimensions.split()[0].split('x'))
+      break
+  proc.wait()
+  assert dimensions, ('Failed to determine video dimensions. output=%s' %
+                      output)
+  return dimensions
+
+
+def _GetFrameTimestampMs(stderr):
+  """Returns the frame timestamp in integer milliseconds from the avconv log.
+
+  The expected line format is:
+  '  dts=1.715  pts=1.715\n'
+
+  We have to be careful to only read a single timestamp per call to avoid
+  deadlock because avconv interleaves its writes to stdout and stderr.
+  """
+  while True:
+    line = ''
+    next_char = ''
+    while next_char != '\n':
+      next_char = stderr.read(1)
+      line += next_char
+    if 'pts=' in line:
+      return int(1000 * float(line.split('=')[-1]))
+
+
+class _Commands(object):
+  """Command types for the external bitmapiter tool."""
+  NEXT = 1
+  PIXELS = 2
+  BOUNDING_BOX = 3
+  HISTOGRAM = 4
+  CROP = 5
+  MAGIC = 0x2751912F
+
+
+class BitmapIter(object):
+  """Extracts Bitmap from video files and supports fast per-pixel operations.
+
+  It communicates with the external bitmapiter tool to process a stream of
+  uncompressed pixels and perform expensive operations: GetBoundingBox and
+  ColorHistogram.
+  """
+
+  def __init__(self, mp4_file):
+    bitmapiter_binary = util.FindSupportBinary('bitmapiter')
+    assert bitmapiter_binary, 'You must build bitmapiter first!'
+
+    self._bpp = 3
+    self._width, self._height = _GetDimensions(mp4_file)
+
+    # Current bitmap, if available.
+    self._bitmap = None
+    # Last timestamp.
+    self._timestamp = None
+
+    # Use rawvideo so that we don't need any external library to parse pixels.
+    command = ['avconv', '-i', mp4_file, '-vcodec', 'rawvideo',
+               '-pix_fmt', 'rgb24', '-dump', '-loglevel', 'debug',
+               '-f', 'rawvideo', '-']
+
+    avconv = subprocess.Popen(command, stdout=subprocess.PIPE,
+                                       stderr=subprocess.PIPE)
+    self._dumpfile = avconv.stderr
+    # Command pipe to communicate to bitmapiter.
+    pipe = os.pipe()
+    bmpit = subprocess.Popen([bitmapiter_binary, str(pipe[0])],
+                             stdin=avconv.stdout,
+                             stdout=subprocess.PIPE,
+                             stderr=subprocess.PIPE)
+    self._cmdpipe = pipe[1]
+    self._outpipe = bmpit.stdout.fileno()
+    self._errfile = bmpit.stderr
+    self._RunCommand(_Commands.MAGIC, self._bpp, self._width, self._height)
+
+  def _RunCommand(self, *command):
+    # Command is a packed bunch of ints.
+    packed_command = struct.pack('i' * len(command), *command)
+    os.write(self._cmdpipe, packed_command)
+    # Response is length + bytes.
+    length_response = os.read(self._outpipe, struct.calcsize('i'))
+    if not length_response:
+      msg = self._errfile.read()
+      if msg == b'STOP\n':
+        raise StopIteration
+      raise Exception(msg)
+
+    length = struct.unpack('i', length_response)[0]
+    if length:
+      return os.read(self._outpipe, length)
+
+  def __iter__(self):
+    return self
+
+  def next(self):
+    self._RunCommand(_Commands.NEXT)
+    self._timestamp = _GetFrameTimestampMs(self._dumpfile)
+    return self
+
+  @property
+  def bpp(self):
+    return self._bpp
+
+  @property
+  def width(self):
+    return self._width
+
+  @property
+  def height(self):
+    return self._height
+
+  @property
+  def timestamp(self):
+    return self._timestamp
+
+  @property
+  def bitmap(self):
+    if not self._bitmap:
+      pixels = self._RunCommand(_Commands.PIXELS)
+      self._bitmap = bitmap.Bitmap(self._bpp, self._width, self._height, pixels)
+    return self._bitmap
+
+  def GetBoundingBox(self, color, tolerance=0):
+    """Finds the minimum box surrounding all occurences of |color|.
+    Returns: (top, left, width, height), match_count
+    Ignores the alpha channel."""
+    response = self._RunCommand(_Commands.BOUNDING_BOX, color, tolerance)
+    top, left, width, height, match_count = struct.unpack('iiiii', response)
+    return (top, left, width, height), match_count
+
+  def Crop(self, left, top, width, height):
+    """Crops the current bitmap down to the specified box."""
+    if (left < 0 or top < 0 or
+        (left + width) > self.width or
+        (top + height) > self.height):
+      raise ValueError('Invalid dimensions')
+    self._RunCommand(_Commands.CROP, left, top, width, height)
+    return self
+
+  def ColorHistogram(self, ignore_color=-1, tolerance=0):
+    """Computes a histogram of the pixel colors in this Bitmap.
+    Args:
+      ignore_color: An RgbaColor to exclude from the bucket counts.
+      tolerance: A tolerance for the ignore_color.
+
+    Returns:
+      A list of 3x256 integers formatted as
+      [r0, r1, ..., g0, g1, ..., b0, b1, ...].
+    """
+    response = self._RunCommand(_Commands.HISTOGRAM, ignore_color, tolerance)
+    out = array.array('i')
+    out.fromstring(response)
+    return out