Adding additional optional parameters to imagediffdb to make calling from Cluster Telemetry possibl…

gm/rebaseline_server/imagediffdb.py

 #!/usr/bin/python
 
 """
 Copyright 2013 Google Inc.
 
 Use of this source code is governed by a BSD-style license that can be
 found in the LICENSE file.
 
 Calulate differences between image pairs, and store them in a database.
 """
 
 import contextlib
 import logging
 import os
 import shutil
 import urllib
 try:
   from PIL import Image, ImageChops
 except ImportError:
   raise ImportError('Requires PIL to be installed; see '
                     + 'http://www.pythonware.com/products/pil/')
 
-IMAGE_SUFFIX = '.png'
+DEFAULT_IMAGE_SUFFIX = '.png'
+DEFAULT_IMAGES_SUBDIR = 'images'
 
-IMAGES_SUBDIR = 'images'
 DIFFS_SUBDIR = 'diffs'
 WHITEDIFFS_SUBDIR = 'whitediffs'
 
 VALUES_PER_BAND = 256
 
 
 class DiffRecord(object):
   """ Record of differences between two images. """
 
   def __init__(self, storage_root,
                expected_image_url, expected_image_locator,
-               actual_image_url, actual_image_locator):
+               actual_image_url, actual_image_locator,
+               expected_images_subdir=DEFAULT_IMAGES_SUBDIR,
+               actual_images_subdir=DEFAULT_IMAGES_SUBDIR,
+               image_suffix=DEFAULT_IMAGE_SUFFIX):
     """Download this pair of images (unless we already have them on local disk),
     and prepare a DiffRecord for them.
 
     TODO(epoger): Make this asynchronously download images, rather than blocking
     until the images have been downloaded and processed.
 
     Args:
       storage_root: root directory on local disk within which we store all
           images
       expected_image_url: file or HTTP url from which we will download the
           expected image
       expected_image_locator: a unique ID string under which we will store the
           expected image within storage_root (probably including a checksum to
           guarantee uniqueness)
       actual_image_url: file or HTTP url from which we will download the
           actual image
       actual_image_locator: a unique ID string under which we will store the
           actual image within storage_root (probably including a checksum to
           guarantee uniqueness)
+      expected_images_subdir: the subdirectory expected images are stored in.
+      actual_images_subdir: the subdirectory actual images are stored in.
+      image_suffix: the suffix of images.
     """
     # Download the expected/actual images, if we don't have them already.
+    # TODO(rmistry): Add a parameter that makes _download_and_open_image raise
+    # an exception if images are not found locally (instead of trying to
+    # download them).
     expected_image = _download_and_open_image(
-        os.path.join(storage_root, IMAGES_SUBDIR,
-                     str(expected_image_locator) + IMAGE_SUFFIX),
+        os.path.join(storage_root, expected_images_subdir,
+                     str(expected_image_locator) + image_suffix),
         expected_image_url)
     actual_image = _download_and_open_image(
-        os.path.join(storage_root, IMAGES_SUBDIR,
-                     str(actual_image_locator) + IMAGE_SUFFIX),
+        os.path.join(storage_root, actual_images_subdir,
+                     str(actual_image_locator) + image_suffix),
         actual_image_url)
 
     # Generate the diff image (absolute diff at each pixel) and
     # max_diff_per_channel.
     diff_image = _generate_image_diff(actual_image, expected_image)
     diff_histogram = diff_image.histogram()
     (diff_width, diff_height) = diff_image.size
     self._weighted_diff_measure = _calculate_weighted_diff_metric(
         diff_histogram, diff_width * diff_height)
     self._max_diff_per_channel = _max_per_band(diff_histogram)
 
     # Generate the whitediff image (any differing pixels show as white).
     # This is tricky, because when you convert color images to grayscale or
     # black & white in PIL, it has its own ideas about thresholds.
     # We have to force it: if a pixel has any color at all, it's a '1'.
     bands = diff_image.split()
     graydiff_image = ImageChops.lighter(ImageChops.lighter(
         bands[0], bands[1]), bands[2])
     whitediff_image = (graydiff_image.point(lambda p: p > 0 and VALUES_PER_BAND)
                                      .convert('1', dither=Image.NONE))
 
     # Final touches on diff_image: use whitediff_image as an alpha mask.
     # Unchanged pixels are transparent; differing pixels are opaque.
     diff_image.putalpha(whitediff_image)
 
     # Store the diff and whitediff images generated above.
     diff_image_locator = _get_difference_locator(
         expected_image_locator=expected_image_locator,
         actual_image_locator=actual_image_locator)
-    basename = str(diff_image_locator) + IMAGE_SUFFIX
+    basename = str(diff_image_locator) + image_suffix
     _save_image(diff_image, os.path.join(
         storage_root, DIFFS_SUBDIR, basename))
     _save_image(whitediff_image, os.path.join(
         storage_root, WHITEDIFFS_SUBDIR, basename))
 
     # Calculate difference metrics.
     (self._width, self._height) = diff_image.size
     self._num_pixels_differing = (
         whitediff_image.histogram()[VALUES_PER_BAND - 1])
 
@@ skipping 212 matching lines @@
   and actual_image.
 
   Args:
     expected_image_locator: locator string pointing at expected image
     actual_image_locator: locator string pointing at actual image
 
   Returns: locator where the diffs between expected and actual images can be
       found
   """
   return "%s-vs-%s" % (expected_image_locator, actual_image_locator)