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