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.
4import base64
5import cStringIO
6
7from telemetry.core import util
8
9util.AddDirToPythonPath(util.GetTelemetryDir(), 'third_party', 'png')
10import png  # pylint: disable=F0401
11
12
13class RgbaColor(object):
14  """Encapsulates an RGBA color retreived from a Bitmap"""
15
16  def __init__(self, r, g, b, a=255):
17    self.r = r
18    self.g = g
19    self.b = b
20    self.a = a
21
22  def IsEqual(self, expected_color, tolerance=0):
23    """Verifies that the color is within a given tolerance of
24    the expected color"""
25    r_diff = abs(self.r - expected_color.r)
26    g_diff = abs(self.g - expected_color.g)
27    b_diff = abs(self.b - expected_color.b)
28    a_diff = abs(self.a - expected_color.a)
29    return (r_diff <= tolerance and g_diff <= tolerance
30        and b_diff <= tolerance and a_diff <= tolerance)
31
32  def AssertIsRGB(self, r, g, b, tolerance=0):
33    assert self.IsEqual(RgbaColor(r, g, b), tolerance)
34
35  def AssertIsRGBA(self, r, g, b, a, tolerance=0):
36    assert self.IsEqual(RgbaColor(r, g, b, a), tolerance)
37
38
39class Bitmap(object):
40  """Utilities for parsing and inspecting a bitmap."""
41
42  def __init__(self, bpp, width, height, pixels, metadata=None):
43    assert bpp in [3, 4], 'Invalid bytes per pixel'
44    assert width > 0, 'Invalid width'
45    assert height > 0, 'Invalid height'
46    assert pixels, 'Must specify pixels'
47    assert bpp * width * height == len(pixels), 'Dimensions and pixels mismatch'
48
49    self._bpp = bpp
50    self._width = width
51    self._height = height
52    self._pixels = pixels
53    self._metadata = metadata or {}
54
55  @property
56  def bpp(self):
57    """Bytes per pixel."""
58    return self._bpp
59
60  @property
61  def width(self):
62    """Width of the bitmap."""
63    return self._width
64
65  @property
66  def height(self):
67    """Height of the bitmap."""
68    return self._height
69
70  @property
71  def pixels(self):
72    """Flat pixel array of the bitmap."""
73    if type(self._pixels) is not bytearray:
74      self._pixels = bytearray(self._pixels)
75    return self._pixels
76
77  @property
78  def metadata(self):
79    self._metadata['size'] = (self.width, self.height)
80    self._metadata['alpha'] = self.bpp == 4
81    self._metadata['bitdepth'] = 8
82    return self._metadata
83
84  def GetPixelColor(self, x, y):
85    """Returns a RgbaColor for the pixel at (x, y)."""
86    base = self._bpp * (y * self._width + x)
87    if self._bpp == 4:
88      return RgbaColor(self._pixels[base + 0], self._pixels[base + 1],
89                       self._pixels[base + 2], self._pixels[base + 3])
90    return RgbaColor(self._pixels[base + 0], self._pixels[base + 1],
91                     self._pixels[base + 2])
92
93  def WritePngFile(self, path):
94    with open(path, "wb") as f:
95      png.Writer(**self.metadata).write_array(f, self.pixels)
96
97  @staticmethod
98  def FromPng(png_data):
99    width, height, pixels, meta = png.Reader(bytes=png_data).read_flat()
100    return Bitmap(4 if meta['alpha'] else 3, width, height, pixels, meta)
101
102  @staticmethod
103  def FromPngFile(path):
104    with open(path, "rb") as f:
105      return Bitmap.FromPng(f.read())
106
107  @staticmethod
108  def FromBase64Png(base64_png):
109    return Bitmap.FromPng(base64.b64decode(base64_png))
110
111  def IsEqual(self, other, tolerance=0):
112    """Determines whether two Bitmaps are identical within a given tolerance."""
113
114    # Dimensions must be equal
115    if self.width != other.width or self.height != other.height:
116      return False
117
118    # Loop over each pixel and test for equality
119    if tolerance or self.bpp != other.bpp:
120      for y in range(self.height):
121        for x in range(self.width):
122          c0 = self.GetPixelColor(x, y)
123          c1 = other.GetPixelColor(x, y)
124          if not c0.IsEqual(c1, tolerance):
125            return False
126    else:
127      return self.pixels == other.pixels
128
129    return True
130
131  def Diff(self, other):
132    """Returns a new Bitmap that represents the difference between this image
133    and another Bitmap."""
134
135    # Output dimensions will be the maximum of the two input dimensions
136    out_width = max(self.width, other.width)
137    out_height = max(self.height, other.height)
138
139    diff = [[0 for x in xrange(out_width * 3)] for x in xrange(out_height)]
140
141    # Loop over each pixel and write out the difference
142    for y in range(out_height):
143      for x in range(out_width):
144        if x < self.width and y < self.height:
145          c0 = self.GetPixelColor(x, y)
146        else:
147          c0 = RgbaColor(0, 0, 0, 0)
148
149        if x < other.width and y < other.height:
150          c1 = other.GetPixelColor(x, y)
151        else:
152          c1 = RgbaColor(0, 0, 0, 0)
153
154        offset = x * 3
155        diff[y][offset] = abs(c0.r - c1.r)
156        diff[y][offset+1] = abs(c0.g - c1.g)
157        diff[y][offset+2] = abs(c0.b - c1.b)
158
159    # This particular method can only save to a file, so the result will be
160    # written into an in-memory buffer and read back into a Bitmap
161    diff_img = png.from_array(diff, mode='RGB')
162    output = cStringIO.StringIO()
163    try:
164      diff_img.save(output)
165      diff = Bitmap.FromPng(output.getvalue())
166    finally:
167      output.close()
168
169    return diff
170
171  def GetBoundingBox(self, color, tolerance=0):
172    """Returns a (top, left, width, height) tuple of the minimum box
173       surrounding all occurences of |color|."""
174    # TODO(szym): Implement this.
175    raise NotImplementedError("GetBoundingBox not yet implemented.")
176
177  def Crop(self, top, left, width, height):
178    """Crops the current bitmap down to the specified box.
179
180    TODO(szym): Make this O(1).
181    """
182    if (left < 0 or top < 0 or
183        (left + width) > self.width or
184        (top + height) > self.height):
185      raise ValueError('Invalid dimensions')
186
187    img_data = [[0 for x in xrange(width * self.bpp)]
188                for y in xrange(height)]
189
190    # Copy each pixel in the sub-rect.
191    # TODO(tonyg): Make this faster by avoiding the copy and artificially
192    # restricting the dimensions.
193    for y in range(height):
194      for x in range(width):
195        c = self.GetPixelColor(x + left, y + top)
196        offset = x * self.bpp
197        img_data[y][offset] = c.r
198        img_data[y][offset + 1] = c.g
199        img_data[y][offset + 2] = c.b
200        if self.bpp == 4:
201          img_data[y][offset + 3] = c.a
202
203    # This particular method can only save to a file, so the result will be
204    # written into an in-memory buffer and read back into a Bitmap
205    crop_img = png.from_array(img_data, mode='RGBA' if self.bpp == 4 else 'RGB')
206    output = cStringIO.StringIO()
207    try:
208      crop_img.save(output)
209      width, height, pixels, meta = png.Reader(
210          bytes=output.getvalue()).read_flat()
211      self._width = width
212      self._height = height
213      self._pixels = pixels
214      self._metadata = meta
215    finally:
216      output.close()
217
218    return self
219
220  def ColorHistogram(self):
221    """Returns a histogram of the pixel colors in this Bitmap."""
222    # TODO(szym): Implement this.
223    raise NotImplementedError("ColorHistogram not yet implemented.")
224