chrome/test/ispy/common/image_tools.py

   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 """Utilities for performing pixel-by-pixel image comparision."""
   6
   7 import itertools
   8 import StringIO
   9 from PIL import Image
  10
  11
  12 def _AreTheSameSize(images):
  13   """Returns whether a set of images are the size size.
  14
  15   Args:
  16     images: a list of images to compare.
  17
  18   Returns:
  19     boolean.
  20
  21   Raises:
  22     Exception: One image or fewer is passed in.
  23   """
  24   if len(images) > 1:
  25     return all(images[0].size == img.size for img in images[1:])
  26   else:
  27     raise Exception('No images passed in.')
  28
  29
  30 def _GetDifferenceWithMask(image1, image2, mask=None,
  31                            masked_color=(225, 225, 225, 255),
  32                            same_color=(255, 255, 255, 255),
  33                            different_color=(210, 0, 0, 255)):
  34   """Returns an image representing the difference between the two images.
  35
  36   This function computes the difference between two images taking into
  37   account a mask if it is provided. The final three arguments represent
  38   the coloration of the generated image.
  39
  40   Args:
  41     image1: the first image to compare.
  42     image2: the second image to compare.
  43     mask: an optional mask image consisting of only black and white pixels
  44       where white pixels indicate the portion of the image to be masked out.
  45     masked_color: the color of a masked section in the resulting image.
  46     same_color: the color of an unmasked section that is the same.
  47       between images 1 and 2 in the resulting image.
  48     different_color: the color of an unmasked section that is different
  49       between images 1 and 2 in the resulting image.
  50
  51   Returns:
  52     A 2-tuple with an image representing the unmasked difference between the
  53     two input images and the number of different pixels.
  54
  55   Raises:
  56     Exception: if image1, image2, and mask are not the same size.
  57   """
  58   image_mask = mask
  59   if not mask:
  60     image_mask = Image.new('RGBA', image1.size, (0, 0, 0, 255))
  61   if not _AreTheSameSize([image1, image2, image_mask]):
  62     raise Exception('images and mask must be the same size.')
  63   image_diff = Image.new('RGBA', image1.size, (0, 0, 0, 255))
  64   data = []
  65   diff_pixels = 0
  66   for m, px1, px2 in itertools.izip(image_mask.getdata(),
  67                                     image1.getdata(),
  68                                     image2.getdata()):
  69     if m == (255, 255, 255, 255):
  70       data.append(masked_color)
  71     elif px1 == px2:
  72       data.append(same_color)
  73     else:
  74       data.append(different_color)
  75       diff_pixels += 1
  76
  77   image_diff.putdata(data)
  78   return (image_diff, diff_pixels)
  79
  80
  81 def CreateMask(images):
  82   """Computes a mask for a set of images.
  83
  84   Returns a difference mask that is computed from the images
  85   which are passed in. The mask will have a white pixel
  86   anywhere that the input images differ and a black pixel
  87   everywhere else.
  88
  89   Args:
  90     images: list of images to compute the mask from.
  91
  92   Returns:
  93     an image of only black and white pixels where white pixels represent
  94       areas in the input images that have differences.
  95
  96   Raises:
  97     Exception: if the images passed in are not of the same size.
  98     Exception: if fewer than one image is passed in.
  99   """
 100   if not images:
 101     raise Exception('mask must be created from one or more images.')
 102   mask = Image.new('RGBA', images[0].size, (0, 0, 0, 255))
 103   image = images[0]
 104   for other_image in images[1:]:
 105     mask = _GetDifferenceWithMask(
 106         image,
 107         other_image,
 108         mask,
 109         masked_color=(255, 255, 255, 255),
 110         same_color=(0, 0, 0, 255),
 111         different_color=(255, 255, 255, 255))[0]
 112   return mask
 113
 114
 115 def AddMasks(masks):
 116   """Combines a list of mask images into one mask image.
 117
 118   Args:
 119     masks: a list of mask-images.
 120
 121   Returns:
 122     a new mask that represents the sum of the masked
 123       regions of the passed in list of mask-images.
 124
 125   Raises:
 126     Exception: if masks is an empty list, or if masks are not the same size.
 127   """
 128   if not masks:
 129     raise Exception('masks must be a list containing at least one image.')
 130   if len(masks) > 1 and not _AreTheSameSize(masks):
 131     raise Exception('masks in list must be of the same size.')
 132   white = (255, 255, 255, 255)
 133   black = (0, 0, 0, 255)
 134   masks_data = [mask.getdata() for mask in masks]
 135   image = Image.new('RGBA', masks[0].size, black)
 136   image.putdata([white if white in px_set else black
 137                  for px_set in itertools.izip(*masks_data)])
 138   return image
 139
 140
 141 def ConvertDiffToMask(diff):
 142   """Converts a Diff image into a Mask image.
 143
 144   Args:
 145     diff: the diff image to convert.
 146
 147   Returns:
 148     a new mask image where everything that was masked or different in the diff
 149     is now masked.
 150   """
 151   white = (255, 255, 255, 255)
 152   black = (0, 0, 0, 255)
 153   diff_data = diff.getdata()
 154   image = Image.new('RGBA', diff.size, black)
 155   image.putdata([black if px == white else white for px in diff_data])
 156   return image
 157
 158
 159 def VisualizeImageDifferences(image1, image2, mask=None):
 160   """Returns an image repesenting the unmasked differences between two images.
 161
 162   Iterates through the pixel values of two images and an optional
 163   mask. If the pixel values are the same, or the pixel is masked,
 164   (0,0,0) is stored for that pixel. Otherwise, (255,255,255) is stored.
 165   This ultimately produces an image where unmasked differences between
 166   the two images are white pixels, and everything else is black.
 167
 168   Args:
 169     image1: an RGB image
 170     image2: another RGB image of the same size as image1.
 171     mask: an optional RGB image consisting of only white and black pixels
 172       where the white pixels represent the parts of the images to be masked
 173       out.
 174
 175   Returns:
 176     A 2-tuple with an image representing the unmasked difference between the
 177     two input images and the number of different pixels.
 178
 179   Raises:
 180     Exception: if the two images and optional mask are different sizes.
 181   """
 182   return _GetDifferenceWithMask(image1, image2, mask)
 183
 184
 185 def InflateMask(image, passes):
 186   """A function that adds layers of pixels around the white edges of a mask.
 187
 188   This function evaluates a 'frontier' of valid pixels indices. Initially,
 189     this frontier contains all indices in the image. However, with each pass
 190     only the pixels' indices which were added to the mask by inflation
 191     are added to the next pass's frontier. This gives the algorithm a
 192     large upfront cost that scales negligably when the number of passes
 193     is increased.
 194
 195   Args:
 196     image: the RGBA PIL.Image mask to inflate.
 197     passes: the number of passes to inflate the image by.
 198
 199   Returns:
 200     A RGBA PIL.Image.
 201   """
 202   inflated = Image.new('RGBA', image.size)
 203   new_dataset = list(image.getdata())
 204   old_dataset = list(image.getdata())
 205
 206   frontier = set(range(len(old_dataset)))
 207   new_frontier = set()
 208
 209   l = [-1, 1]
 210
 211   def _ShadeHorizontal(index, px):
 212     col = index % image.size[0]
 213     if px == (255, 255, 255, 255):
 214       for x in l:
 215         if 0 <= col + x < image.size[0]:
 216           if old_dataset[index + x] != (255, 255, 255, 255):
 217             new_frontier.add(index + x)
 218           new_dataset[index + x] = (255, 255, 255, 255)
 219
 220   def _ShadeVertical(index, px):
 221     row = index / image.size[0]
 222     if px == (255, 255, 255, 255):
 223       for x in l:
 224         if 0 <= row + x < image.size[1]:
 225           if old_dataset[index + image.size[0] * x] != (255, 255, 255, 255):
 226             new_frontier.add(index + image.size[0] * x)
 227           new_dataset[index + image.size[0] * x] = (255, 255, 255, 255)
 228
 229   for _ in range(passes):
 230     for index in frontier:
 231       _ShadeHorizontal(index, old_dataset[index])
 232       _ShadeVertical(index, old_dataset[index])
 233     old_dataset, new_dataset = new_dataset, new_dataset
 234     frontier, new_frontier = new_frontier, set()
 235   inflated.putdata(new_dataset)
 236   return inflated
 237
 238
 239 def TotalDifferentPixels(image1, image2, mask=None):
 240   """Computes the number of different pixels between two images.
 241
 242   Args:
 243     image1: the first RGB image to be compared.
 244     image2: the second RGB image to be compared.
 245     mask: an optional RGB image of only black and white pixels
 246       where white pixels indicate the parts of the image to be masked out.
 247
 248   Returns:
 249     the number of differing pixels between the images.
 250
 251   Raises:
 252     Exception: if the images to be compared and the mask are not the same size.
 253   """
 254   image_mask = mask
 255   if not mask:
 256     image_mask = Image.new('RGBA', image1.size, (0, 0, 0, 255))
 257   if _AreTheSameSize([image1, image2, image_mask]):
 258     total_diff = 0
 259     for px1, px2, m in itertools.izip(image1.getdata(),
 260                                       image2.getdata(),
 261                                       image_mask.getdata()):
 262       if m == (255, 255, 255, 255):
 263         continue
 264       elif px1 != px2:
 265         total_diff += 1
 266       else:
 267         continue
 268     return total_diff
 269   else:
 270     raise Exception('images and mask must be the same size')
 271
 272
 273 def SameImage(image1, image2, mask=None):
 274   """Returns a boolean representing whether the images are the same.
 275
 276   Returns a boolean indicating whether two images are similar
 277   enough to be considered the same. Essentially wraps the
 278   TotalDifferentPixels function.
 279
 280
 281   Args:
 282     image1: an RGB image to compare.
 283     image2: an RGB image to compare.
 284     mask: an optional image of only black and white pixels
 285     where white pixels are masked out
 286
 287   Returns:
 288     True if the images are similar, False otherwise.
 289
 290   Raises:
 291     Exception: if the images (and mask) are different sizes.
 292   """
 293   different_pixels = TotalDifferentPixels(image1, image2, mask)
 294   return different_pixels == 0
 295
 296
 297 def EncodePNG(image):
 298   """Returns the PNG file-contents of the image.
 299
 300   Args:
 301     image: an RGB image to be encoded.
 302
 303   Returns:
 304     a base64 encoded string representing the image.
 305   """
 306   f = StringIO.StringIO()
 307   image.save(f, 'PNG')
 308   encoded_image = f.getvalue()
 309   f.close()
 310   return encoded_image
 311
 312
 313 def DecodePNG(png):
 314   """Returns a RGB image from PNG file-contents.
 315
 316   Args:
 317     encoded_image: PNG file-contents of an RGB image.
 318
 319   Returns:
 320     an RGB image
 321   """
 322   return Image.open(StringIO.StringIO(png))