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"""
6Bitmap is a basic wrapper for image pixels. It includes some basic processing
7tools: crop, find bounding box of a color and compute histogram of color values.
8"""
9
10import array
11import cStringIO
12import struct
13import subprocess
14
15from telemetry.internal.util import binary_manager
16from telemetry.core import platform
17from telemetry.util import color_histogram
18from telemetry.util import rgba_color
19
20import png
21
22
23class _BitmapTools(object):
24  """Wraps a child process of bitmaptools and allows for one command."""
25  CROP_PIXELS = 0
26  HISTOGRAM = 1
27  BOUNDING_BOX = 2
28
29  def __init__(self, dimensions, pixels):
30    binary = binary_manager.FetchPath(
31        'bitmaptools',
32        platform.GetHostPlatform().GetArchName(),
33        platform.GetHostPlatform().GetOSName())
34    assert binary, 'You must build bitmaptools first!'
35
36    self._popen = subprocess.Popen([binary],
37                                   stdin=subprocess.PIPE,
38                                   stdout=subprocess.PIPE,
39                                   stderr=subprocess.PIPE)
40
41    # dimensions are: bpp, width, height, boxleft, boxtop, boxwidth, boxheight
42    packed_dims = struct.pack('iiiiiii', *dimensions)
43    self._popen.stdin.write(packed_dims)
44    # If we got a list of ints, we need to convert it into a byte buffer.
45    if type(pixels) is not bytearray:
46      pixels = bytearray(pixels)
47    self._popen.stdin.write(pixels)
48
49  def _RunCommand(self, *command):
50    assert not self._popen.stdin.closed, (
51      'Exactly one command allowed per instance of tools.')
52    packed_command = struct.pack('i' * len(command), *command)
53    self._popen.stdin.write(packed_command)
54    self._popen.stdin.close()
55    length_packed = self._popen.stdout.read(struct.calcsize('i'))
56    if not length_packed:
57      raise Exception(self._popen.stderr.read())
58    length = struct.unpack('i', length_packed)[0]
59    return self._popen.stdout.read(length)
60
61  def CropPixels(self):
62    return self._RunCommand(_BitmapTools.CROP_PIXELS)
63
64  def Histogram(self, ignore_color, tolerance):
65    ignore_color_int = -1 if ignore_color is None else int(ignore_color)
66    response = self._RunCommand(_BitmapTools.HISTOGRAM,
67                                ignore_color_int, tolerance)
68    out = array.array('i')
69    out.fromstring(response)
70    assert len(out) == 768, (
71        'The ColorHistogram has the wrong number of buckets: %s' % len(out))
72    return color_histogram.ColorHistogram(out[:256], out[256:512], out[512:],
73                                    ignore_color)
74
75  def BoundingBox(self, color, tolerance):
76    response = self._RunCommand(_BitmapTools.BOUNDING_BOX, int(color),
77                                tolerance)
78    unpacked = struct.unpack('iiiii', response)
79    box, count = unpacked[:4], unpacked[-1]
80    if box[2] < 0 or box[3] < 0:
81      box = None
82    return box, count
83
84
85class Bitmap(object):
86  """Utilities for parsing and inspecting a bitmap."""
87
88  def __init__(self, bpp, width, height, pixels, metadata=None):
89    assert bpp in [3, 4], 'Invalid bytes per pixel'
90    assert width > 0, 'Invalid width'
91    assert height > 0, 'Invalid height'
92    assert pixels, 'Must specify pixels'
93    assert bpp * width * height == len(pixels), 'Dimensions and pixels mismatch'
94
95    self._bpp = bpp
96    self._width = width
97    self._height = height
98    self._pixels = pixels
99    self._metadata = metadata or {}
100    self._crop_box = None
101
102  @property
103  def bpp(self):
104    return self._bpp
105
106  @property
107  def width(self):
108    return self._crop_box[2] if self._crop_box else self._width
109
110  @property
111  def height(self):
112    return self._crop_box[3] if self._crop_box else self._height
113
114  def _PrepareTools(self):
115    """Prepares an instance of _BitmapTools which allows exactly one command.
116    """
117    crop_box = self._crop_box or (0, 0, self._width, self._height)
118    return _BitmapTools((self._bpp, self._width, self._height) + crop_box,
119                        self._pixels)
120
121  @property
122  def pixels(self):
123    if self._crop_box:
124      self._pixels = self._PrepareTools().CropPixels()
125      # pylint: disable=unpacking-non-sequence
126      _, _, self._width, self._height = self._crop_box
127      self._crop_box = None
128    if type(self._pixels) is not bytearray:
129      self._pixels = bytearray(self._pixels)
130    return self._pixels
131
132  @property
133  def metadata(self):
134    self._metadata['size'] = (self.width, self.height)
135    self._metadata['alpha'] = self.bpp == 4
136    self._metadata['bitdepth'] = 8
137    return self._metadata
138
139  def GetPixelColor(self, x, y):
140    pixels = self.pixels
141    base = self._bpp * (y * self._width + x)
142    if self._bpp == 4:
143      return rgba_color.RgbaColor(pixels[base + 0], pixels[base + 1],
144                                  pixels[base + 2], pixels[base + 3])
145    return rgba_color.RgbaColor(pixels[base + 0], pixels[base + 1],
146                                pixels[base + 2])
147
148  @staticmethod
149  def FromPng(png_data):
150    width, height, pixels, meta = png.Reader(bytes=png_data).read_flat()
151    return Bitmap(4 if meta['alpha'] else 3, width, height, pixels, meta)
152
153  @staticmethod
154  def FromPngFile(path):
155    with open(path, "rb") as f:
156      return Bitmap.FromPng(f.read())
157
158  def WritePngFile(self, path):
159    with open(path, "wb") as f:
160      png.Writer(**self.metadata).write_array(f, self.pixels)
161
162  def IsEqual(self, other, tolerance=0):
163    # Dimensions must be equal
164    if self.width != other.width or self.height != other.height:
165      return False
166
167    # Loop over each pixel and test for equality
168    if tolerance or self.bpp != other.bpp:
169      for y in range(self.height):
170        for x in range(self.width):
171          c0 = self.GetPixelColor(x, y)
172          c1 = other.GetPixelColor(x, y)
173          if not c0.IsEqual(c1, tolerance):
174            return False
175    else:
176      return self.pixels == other.pixels
177
178    return True
179
180  def Diff(self, other):
181    # Output dimensions will be the maximum of the two input dimensions
182    out_width = max(self.width, other.width)
183    out_height = max(self.height, other.height)
184
185    diff = [[0 for x in xrange(out_width * 3)] for x in xrange(out_height)]
186
187    # Loop over each pixel and write out the difference
188    for y in range(out_height):
189      for x in range(out_width):
190        if x < self.width and y < self.height:
191          c0 = self.GetPixelColor(x, y)
192        else:
193          c0 = rgba_color.RgbaColor(0, 0, 0, 0)
194
195        if x < other.width and y < other.height:
196          c1 = other.GetPixelColor(x, y)
197        else:
198          c1 = rgba_color.RgbaColor(0, 0, 0, 0)
199
200        offset = x * 3
201        diff[y][offset] = abs(c0.r - c1.r)
202        diff[y][offset+1] = abs(c0.g - c1.g)
203        diff[y][offset+2] = abs(c0.b - c1.b)
204
205    # This particular method can only save to a file, so the result will be
206    # written into an in-memory buffer and read back into a Bitmap
207    diff_img = png.from_array(diff, mode='RGB')
208    output = cStringIO.StringIO()
209    try:
210      diff_img.save(output)
211      diff = Bitmap.FromPng(output.getvalue())
212    finally:
213      output.close()
214
215    return diff
216
217  def GetBoundingBox(self, color, tolerance=0):
218    return self._PrepareTools().BoundingBox(color, tolerance)
219
220  def Crop(self, left, top, width, height):
221    cur_box = self._crop_box or (0, 0, self._width, self._height)
222    cur_left, cur_top, cur_width, cur_height = cur_box
223
224    if (left < 0 or top < 0 or
225        (left + width) > cur_width or
226        (top + height) > cur_height):
227      raise ValueError('Invalid dimensions')
228
229    self._crop_box = cur_left + left, cur_top + top, width, height
230    return self
231
232  def ColorHistogram(self, ignore_color=None, tolerance=0):
233    return self._PrepareTools().Histogram(ignore_color, tolerance)
234