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