src/tools/telemetry/telemetry/core/bitmap.py

   1 # Copyright 2013 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 """
   6 Bitmap is a basic wrapper for image pixels. It includes some basic processing
   7 tools: crop, find bounding box of a color and compute histogram of color values.
   8 """
   9
  10 import array
  11 import base64
  12 import cStringIO
  13 import collections
  14 import struct
  15 import subprocess
  16
  17 from telemetry.core import util
  18 from telemetry.core import platform
  19 from telemetry.util import support_binaries
  20
  21 util.AddDirToPythonPath(util.GetTelemetryDir(), 'third_party', 'png')
  22 import png  # pylint: disable=F0401
  23
  24
  25 def HistogramDistance(hist1, hist2):
  26   """Earth mover's distance.
  27
  28   http://en.wikipedia.org/wiki/Earth_mover's_distance
  29   First, normalize the two histograms. Then, treat the two histograms as
  30   piles of dirt, and calculate the cost of turning one pile into the other.
  31
  32   To do this, calculate the difference in one bucket between the two
  33   histograms. Then carry it over in the calculation for the next bucket.
  34   In this way, the difference is weighted by how far it has to move."""
  35   if len(hist1) != len(hist2):
  36     raise ValueError('Trying to compare histograms '
  37       'of different sizes, %s != %s' % (len(hist1), len(hist2)))
  38
  39   n1 = sum(hist1)
  40   n2 = sum(hist2)
  41   if n1 == 0:
  42     raise ValueError('First histogram has 0 pixels in it.')
  43   if n2 == 0:
  44     raise ValueError('Second histogram has 0 pixels in it.')
  45
  46   total = 0
  47   remainder = 0
  48   for value1, value2 in zip(hist1, hist2):
  49     remainder += value1 * n2 - value2 * n1
  50     total += abs(remainder)
  51   assert remainder == 0, (
  52       '%s pixel(s) left over after computing histogram distance.'
  53       % abs(remainder))
  54   return abs(float(total) / n1 / n2)
  55
  56
  57 class ColorHistogram(
  58     collections.namedtuple('ColorHistogram', ['r', 'g', 'b', 'default_color'])):
  59   # pylint: disable=W0232
  60   # pylint: disable=E1002
  61
  62   def __new__(cls, r, g, b, default_color=None):
  63     return super(ColorHistogram, cls).__new__(cls, r, g, b, default_color)
  64
  65   def Distance(self, other):
  66     total = 0
  67     for i in xrange(3):
  68       hist1 = self[i]
  69       hist2 = other[i]
  70
  71       if sum(self[i]) == 0:
  72         if not self.default_color:
  73           raise ValueError('Histogram has no data and no default color.')
  74         hist1 = [0] * 256
  75         hist1[self.default_color[i]] = 1
  76       if sum(other[i]) == 0:
  77         if not other.default_color:
  78           raise ValueError('Histogram has no data and no default color.')
  79         hist2 = [0] * 256
  80         hist2[other.default_color[i]] = 1
  81
  82       total += HistogramDistance(hist1, hist2)
  83     return total
  84
  85
  86 class RgbaColor(collections.namedtuple('RgbaColor', ['r', 'g', 'b', 'a'])):
  87   """Encapsulates an RGBA color retreived from a Bitmap"""
  88   # pylint: disable=W0232
  89   # pylint: disable=E1002
  90
  91   def __new__(cls, r, g, b, a=255):
  92     return super(RgbaColor, cls).__new__(cls, r, g, b, a)
  93
  94   def __int__(self):
  95     return (self.r << 16) | (self.g << 8) | self.b
  96
  97   def IsEqual(self, expected_color, tolerance=0):
  98     """Verifies that the color is within a given tolerance of
  99     the expected color"""
 100     r_diff = abs(self.r - expected_color.r)
 101     g_diff = abs(self.g - expected_color.g)
 102     b_diff = abs(self.b - expected_color.b)
 103     a_diff = abs(self.a - expected_color.a)
 104     return (r_diff <= tolerance and g_diff <= tolerance
 105         and b_diff <= tolerance and a_diff <= tolerance)
 106
 107   def AssertIsRGB(self, r, g, b, tolerance=0):
 108     assert self.IsEqual(RgbaColor(r, g, b), tolerance)
 109
 110   def AssertIsRGBA(self, r, g, b, a, tolerance=0):
 111     assert self.IsEqual(RgbaColor(r, g, b, a), tolerance)
 112
 113
 114 WEB_PAGE_TEST_ORANGE = RgbaColor(222, 100,  13)
 115 WHITE =                RgbaColor(255, 255, 255)
 116
 117
 118 class _BitmapTools(object):
 119   """Wraps a child process of bitmaptools and allows for one command."""
 120   CROP_PIXELS = 0
 121   HISTOGRAM = 1
 122   BOUNDING_BOX = 2
 123
 124   def __init__(self, dimensions, pixels):
 125     binary = support_binaries.FindPath(
 126         'bitmaptools',
 127         platform.GetHostPlatform().GetArchName(),
 128         platform.GetHostPlatform().GetOSName())
 129     assert binary, 'You must build bitmaptools first!'
 130
 131     self._popen = subprocess.Popen([binary],
 132                                    stdin=subprocess.PIPE,
 133                                    stdout=subprocess.PIPE,
 134                                    stderr=subprocess.PIPE)
 135
 136     # dimensions are: bpp, width, height, boxleft, boxtop, boxwidth, boxheight
 137     packed_dims = struct.pack('iiiiiii', *dimensions)
 138     self._popen.stdin.write(packed_dims)
 139     # If we got a list of ints, we need to convert it into a byte buffer.
 140     if type(pixels) is not bytearray:
 141       pixels = bytearray(pixels)
 142     self._popen.stdin.write(pixels)
 143
 144   def _RunCommand(self, *command):
 145     assert not self._popen.stdin.closed, (
 146       'Exactly one command allowed per instance of tools.')
 147     packed_command = struct.pack('i' * len(command), *command)
 148     self._popen.stdin.write(packed_command)
 149     self._popen.stdin.close()
 150     length_packed = self._popen.stdout.read(struct.calcsize('i'))
 151     if not length_packed:
 152       raise Exception(self._popen.stderr.read())
 153     length = struct.unpack('i', length_packed)[0]
 154     return self._popen.stdout.read(length)
 155
 156   def CropPixels(self):
 157     return self._RunCommand(_BitmapTools.CROP_PIXELS)
 158
 159   def Histogram(self, ignore_color, tolerance):
 160     ignore_color_int = -1 if ignore_color is None else int(ignore_color)
 161     response = self._RunCommand(_BitmapTools.HISTOGRAM,
 162                                 ignore_color_int, tolerance)
 163     out = array.array('i')
 164     out.fromstring(response)
 165     assert len(out) == 768, (
 166         'The ColorHistogram has the wrong number of buckets: %s' % len(out))
 167     return ColorHistogram(out[:256], out[256:512], out[512:], ignore_color)
 168
 169   def BoundingBox(self, color, tolerance):
 170     response = self._RunCommand(_BitmapTools.BOUNDING_BOX, int(color),
 171                                 tolerance)
 172     unpacked = struct.unpack('iiiii', response)
 173     box, count = unpacked[:4], unpacked[-1]
 174     if box[2] < 0 or box[3] < 0:
 175       box = None
 176     return box, count
 177
 178
 179 class Bitmap(object):
 180   """Utilities for parsing and inspecting a bitmap."""
 181
 182   def __init__(self, bpp, width, height, pixels, metadata=None):
 183     assert bpp in [3, 4], 'Invalid bytes per pixel'
 184     assert width > 0, 'Invalid width'
 185     assert height > 0, 'Invalid height'
 186     assert pixels, 'Must specify pixels'
 187     assert bpp * width * height == len(pixels), 'Dimensions and pixels mismatch'
 188
 189     self._bpp = bpp
 190     self._width = width
 191     self._height = height
 192     self._pixels = pixels
 193     self._metadata = metadata or {}
 194     self._crop_box = None
 195
 196   @property
 197   def bpp(self):
 198     """Bytes per pixel."""
 199     return self._bpp
 200
 201   @property
 202   def width(self):
 203     """Width of the bitmap."""
 204     return self._crop_box[2] if self._crop_box else self._width
 205
 206   @property
 207   def height(self):
 208     """Height of the bitmap."""
 209     return self._crop_box[3] if self._crop_box else self._height
 210
 211   def _PrepareTools(self):
 212     """Prepares an instance of _BitmapTools which allows exactly one command.
 213     """
 214     crop_box = self._crop_box or (0, 0, self._width, self._height)
 215     return _BitmapTools((self._bpp, self._width, self._height) + crop_box,
 216                         self._pixels)
 217
 218   @property
 219   def pixels(self):
 220     """Flat pixel array of the bitmap."""
 221     if self._crop_box:
 222       self._pixels = self._PrepareTools().CropPixels()
 223       _, _, self._width, self._height = self._crop_box
 224       self._crop_box = None
 225     if type(self._pixels) is not bytearray:
 226       self._pixels = bytearray(self._pixels)
 227     return self._pixels
 228
 229   @property
 230   def metadata(self):
 231     self._metadata['size'] = (self.width, self.height)
 232     self._metadata['alpha'] = self.bpp == 4
 233     self._metadata['bitdepth'] = 8
 234     return self._metadata
 235
 236   def GetPixelColor(self, x, y):
 237     """Returns a RgbaColor for the pixel at (x, y)."""
 238     pixels = self.pixels
 239     base = self._bpp * (y * self._width + x)
 240     if self._bpp == 4:
 241       return RgbaColor(pixels[base + 0], pixels[base + 1],
 242                        pixels[base + 2], pixels[base + 3])
 243     return RgbaColor(pixels[base + 0], pixels[base + 1],
 244                      pixels[base + 2])
 245
 246   def WritePngFile(self, path):
 247     with open(path, "wb") as f:
 248       png.Writer(**self.metadata).write_array(f, self.pixels)
 249
 250   @staticmethod
 251   def FromPng(png_data):
 252     width, height, pixels, meta = png.Reader(bytes=png_data).read_flat()
 253     return Bitmap(4 if meta['alpha'] else 3, width, height, pixels, meta)
 254
 255   @staticmethod
 256   def FromPngFile(path):
 257     with open(path, "rb") as f:
 258       return Bitmap.FromPng(f.read())
 259
 260   @staticmethod
 261   def FromBase64Png(base64_png):
 262     return Bitmap.FromPng(base64.b64decode(base64_png))
 263
 264   def IsEqual(self, other, tolerance=0):
 265     """Determines whether two Bitmaps are identical within a given tolerance."""
 266
 267     # Dimensions must be equal
 268     if self.width != other.width or self.height != other.height:
 269       return False
 270
 271     # Loop over each pixel and test for equality
 272     if tolerance or self.bpp != other.bpp:
 273       for y in range(self.height):
 274         for x in range(self.width):
 275           c0 = self.GetPixelColor(x, y)
 276           c1 = other.GetPixelColor(x, y)
 277           if not c0.IsEqual(c1, tolerance):
 278             return False
 279     else:
 280       return self.pixels == other.pixels
 281
 282     return True
 283
 284   def Diff(self, other):
 285     """Returns a new Bitmap that represents the difference between this image
 286     and another Bitmap."""
 287
 288     # Output dimensions will be the maximum of the two input dimensions
 289     out_width = max(self.width, other.width)
 290     out_height = max(self.height, other.height)
 291
 292     diff = [[0 for x in xrange(out_width * 3)] for x in xrange(out_height)]
 293
 294     # Loop over each pixel and write out the difference
 295     for y in range(out_height):
 296       for x in range(out_width):
 297         if x < self.width and y < self.height:
 298           c0 = self.GetPixelColor(x, y)
 299         else:
 300           c0 = RgbaColor(0, 0, 0, 0)
 301
 302         if x < other.width and y < other.height:
 303           c1 = other.GetPixelColor(x, y)
 304         else:
 305           c1 = RgbaColor(0, 0, 0, 0)
 306
 307         offset = x * 3
 308         diff[y][offset] = abs(c0.r - c1.r)
 309         diff[y][offset+1] = abs(c0.g - c1.g)
 310         diff[y][offset+2] = abs(c0.b - c1.b)
 311
 312     # This particular method can only save to a file, so the result will be
 313     # written into an in-memory buffer and read back into a Bitmap
 314     diff_img = png.from_array(diff, mode='RGB')
 315     output = cStringIO.StringIO()
 316     try:
 317       diff_img.save(output)
 318       diff = Bitmap.FromPng(output.getvalue())
 319     finally:
 320       output.close()
 321
 322     return diff
 323
 324   def GetBoundingBox(self, color, tolerance=0):
 325     """Finds the minimum box surrounding all occurences of |color|.
 326     Returns: (top, left, width, height), match_count
 327     Ignores the alpha channel."""
 328     return self._PrepareTools().BoundingBox(color, tolerance)
 329
 330   def Crop(self, left, top, width, height):
 331     """Crops the current bitmap down to the specified box."""
 332     cur_box = self._crop_box or (0, 0, self._width, self._height)
 333     cur_left, cur_top, cur_width, cur_height = cur_box
 334
 335     if (left < 0 or top < 0 or
 336         (left + width) > cur_width or
 337         (top + height) > cur_height):
 338       raise ValueError('Invalid dimensions')
 339
 340     self._crop_box = cur_left + left, cur_top + top, width, height
 341     return self
 342
 343   def ColorHistogram(self, ignore_color=None, tolerance=0):
 344     """Computes a histogram of the pixel colors in this Bitmap.
 345     Args:
 346       ignore_color: An RgbaColor to exclude from the bucket counts.
 347       tolerance: A tolerance for the ignore_color.
 348
 349     Returns:
 350       A ColorHistogram namedtuple with 256 integers in each field: r, g, and b.
 351     """
 352     return self._PrepareTools().Histogram(ignore_color, tolerance)