| # Copyright 2013 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. |
| |
| """ |
| Bitmap is a basic wrapper for image pixels. It includes some basic processing |
| tools: crop, find bounding box of a color and compute histogram of color values. |
| """ |
| |
| import array |
| import base64 |
| import cStringIO |
| import collections |
| import struct |
| import subprocess |
| |
| from telemetry.core import util |
| from telemetry.core import platform |
| from telemetry.util import support_binaries |
| |
| util.AddDirToPythonPath(util.GetTelemetryDir(), 'third_party', 'png') |
| import png # pylint: disable=F0401 |
| |
| |
| def HistogramDistance(hist1, hist2): |
| """Earth mover's distance. |
| |
| http://en.wikipedia.org/wiki/Earth_mover's_distance |
| First, normalize the two histograms. Then, treat the two histograms as |
| piles of dirt, and calculate the cost of turning one pile into the other. |
| |
| To do this, calculate the difference in one bucket between the two |
| histograms. Then carry it over in the calculation for the next bucket. |
| In this way, the difference is weighted by how far it has to move.""" |
| if len(hist1) != len(hist2): |
| raise ValueError('Trying to compare histograms ' |
| 'of different sizes, %s != %s' % (len(hist1), len(hist2))) |
| |
| n1 = sum(hist1) |
| n2 = sum(hist2) |
| if n1 == 0: |
| raise ValueError('First histogram has 0 pixels in it.') |
| if n2 == 0: |
| raise ValueError('Second histogram has 0 pixels in it.') |
| |
| total = 0 |
| remainder = 0 |
| for value1, value2 in zip(hist1, hist2): |
| remainder += value1 * n2 - value2 * n1 |
| total += abs(remainder) |
| assert remainder == 0, ( |
| '%s pixel(s) left over after computing histogram distance.' |
| % abs(remainder)) |
| return abs(float(total) / n1 / n2) |
| |
| |
| class ColorHistogram( |
| collections.namedtuple('ColorHistogram', ['r', 'g', 'b', 'default_color'])): |
| # pylint: disable=W0232 |
| # pylint: disable=E1002 |
| |
| def __new__(cls, r, g, b, default_color=None): |
| return super(ColorHistogram, cls).__new__(cls, r, g, b, default_color) |
| |
| def Distance(self, other): |
| total = 0 |
| for i in xrange(3): |
| hist1 = self[i] |
| hist2 = other[i] |
| |
| if sum(self[i]) == 0: |
| if not self.default_color: |
| raise ValueError('Histogram has no data and no default color.') |
| hist1 = [0] * 256 |
| hist1[self.default_color[i]] = 1 |
| if sum(other[i]) == 0: |
| if not other.default_color: |
| raise ValueError('Histogram has no data and no default color.') |
| hist2 = [0] * 256 |
| hist2[other.default_color[i]] = 1 |
| |
| total += HistogramDistance(hist1, hist2) |
| return total |
| |
| |
| class RgbaColor(collections.namedtuple('RgbaColor', ['r', 'g', 'b', 'a'])): |
| """Encapsulates an RGBA color retreived from a Bitmap""" |
| # pylint: disable=W0232 |
| # pylint: disable=E1002 |
| |
| def __new__(cls, r, g, b, a=255): |
| return super(RgbaColor, cls).__new__(cls, r, g, b, a) |
| |
| def __int__(self): |
| return (self.r << 16) | (self.g << 8) | self.b |
| |
| def IsEqual(self, expected_color, tolerance=0): |
| """Verifies that the color is within a given tolerance of |
| the expected color""" |
| r_diff = abs(self.r - expected_color.r) |
| g_diff = abs(self.g - expected_color.g) |
| b_diff = abs(self.b - expected_color.b) |
| a_diff = abs(self.a - expected_color.a) |
| return (r_diff <= tolerance and g_diff <= tolerance |
| and b_diff <= tolerance and a_diff <= tolerance) |
| |
| def AssertIsRGB(self, r, g, b, tolerance=0): |
| assert self.IsEqual(RgbaColor(r, g, b), tolerance) |
| |
| def AssertIsRGBA(self, r, g, b, a, tolerance=0): |
| assert self.IsEqual(RgbaColor(r, g, b, a), tolerance) |
| |
| |
| WEB_PAGE_TEST_ORANGE = RgbaColor(222, 100, 13) |
| WHITE = RgbaColor(255, 255, 255) |
| |
| |
| class _BitmapTools(object): |
| """Wraps a child process of bitmaptools and allows for one command.""" |
| CROP_PIXELS = 0 |
| HISTOGRAM = 1 |
| BOUNDING_BOX = 2 |
| |
| def __init__(self, dimensions, pixels): |
| binary = support_binaries.FindPath( |
| 'bitmaptools', |
| platform.GetHostPlatform().GetArchName(), |
| platform.GetHostPlatform().GetOSName()) |
| assert binary, 'You must build bitmaptools first!' |
| |
| self._popen = subprocess.Popen([binary], |
| stdin=subprocess.PIPE, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE) |
| |
| # dimensions are: bpp, width, height, boxleft, boxtop, boxwidth, boxheight |
| packed_dims = struct.pack('iiiiiii', *dimensions) |
| self._popen.stdin.write(packed_dims) |
| # If we got a list of ints, we need to convert it into a byte buffer. |
| if type(pixels) is not bytearray: |
| pixels = bytearray(pixels) |
| self._popen.stdin.write(pixels) |
| |
| def _RunCommand(self, *command): |
| assert not self._popen.stdin.closed, ( |
| 'Exactly one command allowed per instance of tools.') |
| packed_command = struct.pack('i' * len(command), *command) |
| self._popen.stdin.write(packed_command) |
| self._popen.stdin.close() |
| length_packed = self._popen.stdout.read(struct.calcsize('i')) |
| if not length_packed: |
| raise Exception(self._popen.stderr.read()) |
| length = struct.unpack('i', length_packed)[0] |
| return self._popen.stdout.read(length) |
| |
| def CropPixels(self): |
| return self._RunCommand(_BitmapTools.CROP_PIXELS) |
| |
| def Histogram(self, ignore_color, tolerance): |
| ignore_color_int = -1 if ignore_color is None else int(ignore_color) |
| response = self._RunCommand(_BitmapTools.HISTOGRAM, |
| ignore_color_int, tolerance) |
| out = array.array('i') |
| out.fromstring(response) |
| assert len(out) == 768, ( |
| 'The ColorHistogram has the wrong number of buckets: %s' % len(out)) |
| return ColorHistogram(out[:256], out[256:512], out[512:], ignore_color) |
| |
| def BoundingBox(self, color, tolerance): |
| response = self._RunCommand(_BitmapTools.BOUNDING_BOX, int(color), |
| tolerance) |
| unpacked = struct.unpack('iiiii', response) |
| box, count = unpacked[:4], unpacked[-1] |
| if box[2] < 0 or box[3] < 0: |
| box = None |
| return box, count |
| |
| |
| class Bitmap(object): |
| """Utilities for parsing and inspecting a bitmap.""" |
| |
| def __init__(self, bpp, width, height, pixels, metadata=None): |
| assert bpp in [3, 4], 'Invalid bytes per pixel' |
| assert width > 0, 'Invalid width' |
| assert height > 0, 'Invalid height' |
| assert pixels, 'Must specify pixels' |
| assert bpp * width * height == len(pixels), 'Dimensions and pixels mismatch' |
| |
| self._bpp = bpp |
| self._width = width |
| self._height = height |
| self._pixels = pixels |
| self._metadata = metadata or {} |
| self._crop_box = None |
| |
| @property |
| def bpp(self): |
| """Bytes per pixel.""" |
| return self._bpp |
| |
| @property |
| def width(self): |
| """Width of the bitmap.""" |
| return self._crop_box[2] if self._crop_box else self._width |
| |
| @property |
| def height(self): |
| """Height of the bitmap.""" |
| return self._crop_box[3] if self._crop_box else self._height |
| |
| def _PrepareTools(self): |
| """Prepares an instance of _BitmapTools which allows exactly one command. |
| """ |
| crop_box = self._crop_box or (0, 0, self._width, self._height) |
| return _BitmapTools((self._bpp, self._width, self._height) + crop_box, |
| self._pixels) |
| |
| @property |
| def pixels(self): |
| """Flat pixel array of the bitmap.""" |
| if self._crop_box: |
| self._pixels = self._PrepareTools().CropPixels() |
| _, _, self._width, self._height = self._crop_box |
| self._crop_box = None |
| if type(self._pixels) is not bytearray: |
| self._pixels = bytearray(self._pixels) |
| return self._pixels |
| |
| @property |
| def metadata(self): |
| self._metadata['size'] = (self.width, self.height) |
| self._metadata['alpha'] = self.bpp == 4 |
| self._metadata['bitdepth'] = 8 |
| return self._metadata |
| |
| def GetPixelColor(self, x, y): |
| """Returns a RgbaColor for the pixel at (x, y).""" |
| pixels = self.pixels |
| base = self._bpp * (y * self._width + x) |
| if self._bpp == 4: |
| return RgbaColor(pixels[base + 0], pixels[base + 1], |
| pixels[base + 2], pixels[base + 3]) |
| return RgbaColor(pixels[base + 0], pixels[base + 1], |
| pixels[base + 2]) |
| |
| def WritePngFile(self, path): |
| with open(path, "wb") as f: |
| png.Writer(**self.metadata).write_array(f, self.pixels) |
| |
| @staticmethod |
| def FromPng(png_data): |
| width, height, pixels, meta = png.Reader(bytes=png_data).read_flat() |
| return Bitmap(4 if meta['alpha'] else 3, width, height, pixels, meta) |
| |
| @staticmethod |
| def FromPngFile(path): |
| with open(path, "rb") as f: |
| return Bitmap.FromPng(f.read()) |
| |
| @staticmethod |
| def FromBase64Png(base64_png): |
| return Bitmap.FromPng(base64.b64decode(base64_png)) |
| |
| def IsEqual(self, other, tolerance=0): |
| """Determines whether two Bitmaps are identical within a given tolerance.""" |
| |
| # Dimensions must be equal |
| if self.width != other.width or self.height != other.height: |
| return False |
| |
| # Loop over each pixel and test for equality |
| if tolerance or self.bpp != other.bpp: |
| for y in range(self.height): |
| for x in range(self.width): |
| c0 = self.GetPixelColor(x, y) |
| c1 = other.GetPixelColor(x, y) |
| if not c0.IsEqual(c1, tolerance): |
| return False |
| else: |
| return self.pixels == other.pixels |
| |
| return True |
| |
| def Diff(self, other): |
| """Returns a new Bitmap that represents the difference between this image |
| and another Bitmap.""" |
| |
| # Output dimensions will be the maximum of the two input dimensions |
| out_width = max(self.width, other.width) |
| out_height = max(self.height, other.height) |
| |
| diff = [[0 for x in xrange(out_width * 3)] for x in xrange(out_height)] |
| |
| # Loop over each pixel and write out the difference |
| for y in range(out_height): |
| for x in range(out_width): |
| if x < self.width and y < self.height: |
| c0 = self.GetPixelColor(x, y) |
| else: |
| c0 = RgbaColor(0, 0, 0, 0) |
| |
| if x < other.width and y < other.height: |
| c1 = other.GetPixelColor(x, y) |
| else: |
| c1 = RgbaColor(0, 0, 0, 0) |
| |
| offset = x * 3 |
| diff[y][offset] = abs(c0.r - c1.r) |
| diff[y][offset+1] = abs(c0.g - c1.g) |
| diff[y][offset+2] = abs(c0.b - c1.b) |
| |
| # This particular method can only save to a file, so the result will be |
| # written into an in-memory buffer and read back into a Bitmap |
| diff_img = png.from_array(diff, mode='RGB') |
| output = cStringIO.StringIO() |
| try: |
| diff_img.save(output) |
| diff = Bitmap.FromPng(output.getvalue()) |
| finally: |
| output.close() |
| |
| return diff |
| |
| 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.""" |
| return self._PrepareTools().BoundingBox(color, tolerance) |
| |
| def Crop(self, left, top, width, height): |
| """Crops the current bitmap down to the specified box.""" |
| cur_box = self._crop_box or (0, 0, self._width, self._height) |
| cur_left, cur_top, cur_width, cur_height = cur_box |
| |
| if (left < 0 or top < 0 or |
| (left + width) > cur_width or |
| (top + height) > cur_height): |
| raise ValueError('Invalid dimensions') |
| |
| self._crop_box = cur_left + left, cur_top + top, width, height |
| return self |
| |
| def ColorHistogram(self, ignore_color=None, 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 ColorHistogram namedtuple with 256 integers in each field: r, g, and b. |
| """ |
| return self._PrepareTools().Histogram(ignore_color, tolerance) |