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
7import itertools
8import StringIO
9from PIL import Image
10
11
12def _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
30def _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
81def 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
115def 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
141def 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
159def 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
185def 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
239def 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
273def 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
297def 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
313def 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))
323