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