Implement 3 missing Imagenet-C corruptions. Resize and crop images before applying corruptions. Save corrupted images as JPEG files. Change corruption names to reflect names in original release.

PiperOrigin-RevId: 289176965

Author: normster

tensorflow_datasets/image/corruptions.py

@@ -15,22 +15,33 @@
 
 """Common corruptions to images.
 
-Define 12+4 common image corruptions: Gaussian noise, shot noise, impulse_noise,
+Define 15+4 common image corruptions: Gaussian noise, shot noise, impulse_noise,
 defocus blur, frosted glass blur, zoom blur, fog, brightness, contrast, elastic,
-pixelate, jpeg compression.
+pixelate, jpeg compression, frost, snow, and motion blur.
 
-4 extra corruptions include gaussian blur, saturate, spatter, and speckle
-noise.
+4 extra corruptions: gaussian blur, saturate, spatter, and speckle noise.
 """
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
 import io
+import subprocess
+import tempfile
 import numpy as np
+import tensorflow.compat.v2 as tf
 import tensorflow_datasets.public_api as tfds
 
+
+# To be populated by download_manager
+FROST_FILENAMES = []
+
+
+def _imagemagick_bin():
+  return 'imagemagick'  # pylint: disable=unreachable
+
+
 # /////////////// Corruption Helpers ///////////////
 
 
@@ -239,7 +250,7 @@ def defocus_blur(x, severity=1):
   return around_and_astype(x_clip)
 
 
-def frosted_glass_blur(x, severity=1):
+def glass_blur(x, severity=1):
   """Frosted glass blurring to images.
 
   Apply frosted glass blurring to images by shuffling pixels locally.
@@ -367,7 +378,7 @@ def contrast(x, severity=1):
   return around_and_astype(x_clip)
 
 
-def elastic(x, severity=1):
+def elastic_transform(x, severity=1):
   """Conduct elastic transform to images.
 
   Elastic transform is performed on small patches of the images.
@@ -469,6 +480,131 @@ def jpeg_compression(x, severity=1):
   return np.asarray(x)
 
 
+def frost(x, severity=1):
+  """Apply frost to images.
+
+  Args:
+    x: numpy array, uncorrupted image, assumed to have uint8 pixel in [0,255].
+    severity: integer, severity of corruption.
+
+  Returns:
+    numpy array, image with uint8 pixels in [0,255]. Applied frost.
+  """
+  c = [(1, 0.4), (0.8, 0.6), (0.7, 0.7), (0.65, 0.7), (0.6, 0.75)][severity - 1]
+  filename = FROST_FILENAMES[np.random.randint(5)]
+  with tempfile.NamedTemporaryFile() as im_frost:
+    tf.io.gfile.copy(filename, im_frost.name, overwrite=True)
+    frost_img = tfds.core.lazy_imports.cv2.imread(im_frost.name)
+  # randomly crop and convert to rgb
+  x_start, y_start = np.random.randint(
+      0, frost_img.shape[0] - 224), np.random.randint(0,
+                                                      frost_img.shape[1] - 224)
+  frost_img = frost_img[x_start:x_start + 224, y_start:y_start + 224][...,
+                                                                      [2, 1, 0]]
+
+  x = np.clip(c[0] * np.array(x) + c[1] * frost_img, 0, 255)
+
+  return around_and_astype(x)
+
+
+def snow(x, severity=1):
+  """Apply snow to images.
+
+  Args:
+    x: numpy array, uncorrupted image, assumed to have uint8 pixel in [0,255].
+    severity: integer, severity of corruption.
+
+  Returns:
+    numpy array, image with uint8 pixels in [0,255]. Applied snow.
+  """
+  cv2 = tfds.core.lazy_imports.cv2
+  PIL_Image = tfds.core.lazy_imports.PIL_Image  # pylint: disable=invalid-name
+  c = [(0.1, 0.3, 3, 0.5, 10, 4, 0.8), (0.2, 0.3, 2, 0.5, 12, 4, 0.7),
+       (0.55, 0.3, 4, 0.9, 12, 8, 0.7), (0.55, 0.3, 4.5, 0.85, 12, 8, 0.65),
+       (0.55, 0.3, 2.5, 0.85, 12, 12, 0.55)][severity - 1]
+
+  x = np.array(x, dtype=np.float32) / 255.
+  snow_layer = np.random.normal(
+      size=x.shape[:2], loc=c[0], scale=c[1])  # [:2] for monochrome
+
+  snow_layer = clipped_zoom(snow_layer[..., np.newaxis], c[2])
+  snow_layer[snow_layer < c[3]] = 0
+
+  snow_layer = PIL_Image.fromarray(
+      (np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode='L')
+
+  with tempfile.NamedTemporaryFile() as im_input:
+    with tempfile.NamedTemporaryFile() as im_output:
+      snow_layer.save(im_input.name, format='PNG')
+
+      convert_bin = _imagemagick_bin()
+      radius = c[4]
+      sigma = c[5]
+      angle = np.random.uniform(-135, -45)
+
+      subprocess.check_output([
+          convert_bin, '-motion-blur', '{}x{}+{}'.format(radius, sigma, angle),
+          im_input.name, im_output.name
+      ])
+
+      with open(im_output.name, 'rb') as f:
+        output = f.read()
+
+  snow_layer = cv2.imdecode(
+      np.fromstring(output, np.uint8), cv2.IMREAD_UNCHANGED) / 255.
+  snow_layer = snow_layer[..., np.newaxis]
+
+  x = c[6] * x + (1 - c[6]) * np.maximum(
+      x,
+      cv2.cvtColor(x, cv2.COLOR_RGB2GRAY).reshape(224, 224, 1) * 1.5 + 0.5)
+  x = np.clip(x + snow_layer + np.rot90(snow_layer, k=2), 0, 1) * 255
+
+  return around_and_astype(x)
+
+
+def motion_blur(x, severity=1):
+  """Apply motion blur to images.
+
+  Args:
+    x: numpy array, uncorrupted image, assumed to have uint8 pixel in [0,255].
+    severity: integer, severity of corruption.
+
+  Returns:
+    numpy array, image with uint8 pixels in [0,255]. Applied motion blur.
+  """
+  c = [(10, 3), (15, 5), (15, 8), (15, 12), (20, 15)][severity - 1]
+
+  x = tfds.core.lazy_imports.PIL_Image.fromarray(x.astype(np.uint8))
+
+  with tempfile.NamedTemporaryFile() as im_input:
+    with tempfile.NamedTemporaryFile() as im_output:
+      x.save(im_input.name, format='PNG')
+
+      convert_bin = _imagemagick_bin()
+      radius = c[0]
+      sigma = c[1]
+      angle = np.random.uniform(-45, -45)
+
+      subprocess.check_output([
+          convert_bin, '-motion-blur', '{}x{}+{}'.format(radius, sigma, angle),
+          im_input.name, im_output.name
+      ])
+
+      with open(im_output.name, 'rb') as f:
+        output = f.read()
+
+  x = tfds.core.lazy_imports.cv2.imdecode(
+      np.fromstring(output, np.uint8),
+      tfds.core.lazy_imports.cv2.IMREAD_UNCHANGED)
+
+  if x.shape != (224, 224):
+    x = np.clip(x[..., [2, 1, 0]], 0, 255)  # BGR to RGB
+  else:  # greyscale to RGB
+    x = np.clip(np.array([x, x, x]).transpose((1, 2, 0)), 0, 255)
+
+  return around_and_astype(x)
+
+
 # /////////////// Extra Corruptions ///////////////
 
 
@@ -484,7 +620,7 @@ def gaussian_blur(x, severity=1):
   """
   c = [1, 2, 3, 4, 6][severity - 1]
 
-  x = tfds.core.lazy_imports.gaussian(
+  x = tfds.core.lazy_imports.skimage.filters.gaussian(
       np.array(x) / 255., sigma=c, multichannel=True)
   x = np.clip(x, 0, 1) * 255
 
@@ -543,7 +679,7 @@ def spatter(x, severity=1):
     #     ker = np.array([[-1,-2,-3],[-2,0,0],[-3,0,1]], dtype=np.float32)
     #     ker -= np.mean(ker)
     ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
-    dist = cv2.filter2D(dist, cv2.CV_8U, ker)
+    dist = cv2.filter2D(dist, cv2.CVX_8U, ker)
     dist = cv2.blur(dist, (3, 3)).astype(np.float32)
 
     m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)

tensorflow_datasets/image/imagenet2012_corrupted.py

@@ -30,10 +30,7 @@
 _DESCRIPTION = """\
 Imagenet2012Corrupted is a dataset generated by adding common corruptions to the
 images in the ImageNet dataset. In the original paper, there are 15 + 4
-different corruptions, and each has 5 levels of severity. In this dataset, we
-implement 12 out of the 15 corruptions, including Gaussian noise, shot noise,
-impulse_noise, defocus blur, frosted glass blur, zoom blur, fog, brightness,
-contrast, elastic, pixelate, and jpeg compression. We also implement the 4 extra
+different corruptions, and each has 5 levels of severity. We also implement the 4 extra
 corruptions gaussian blur, saturate, spatter, and speckle noise. The randomness
 is fixed so that regeneration is deterministic.
 """
@@ -56,25 +53,37 @@
 # tar file).
 _VALIDATION_LABELS_FNAME = 'image/imagenet2012_validation_labels.txt'
 
-# TODO(normanmu): implement frost, snow, and motion blur once wand library is
-# upgraded (cl/262186801)
+_FROST_FILEBASE = 'https://raw.githubusercontent.com/hendrycks/robustness/master/ImageNet-C/imagenet_c/imagenet_c/frost/'
+_FROST_FILENAMES = [
+    _FROST_FILEBASE + f for f in [
+        'frost1.png', 'frost2.png', 'frost3.png', 'frost4.jpg', 'frost5.jpg',
+        'frost6.jpg'
+    ]
+]
+
 BENCHMARK_CORRUPTIONS = [
     'gaussian_noise',
     'shot_noise',
     'impulse_noise',
     'defocus_blur',
-    'frosted_glass_blur',
+    'glass_blur',
+    'motion_blur',
     'zoom_blur',
+    'snow',
+    'frost',
     'fog',
     'brightness',
     'contrast',
-    'elastic',
+    'elastic_transform',
     'pixelate',
     'jpeg_compression',
 ]
 
 EXTRA_CORRUPTIONS = ['gaussian_blur', 'saturate', 'spatter', 'speckle_noise']
 
+_IMAGE_SIZE = 224
+_CROP_PADDING = 32
+
 
 class Imagenet2012CorruptedConfig(tfds.core.BuilderConfig):
   """BuilderConfig for Imagenet2012Corrupted."""
@@ -94,23 +103,21 @@ def __init__(self, corruption_type=None, severity=1, **kwargs):
     self.severity = severity
 
 
-_VERSION = tfds.core.Version(
-    '0.0.1', experiments={tfds.core.Experiment.S3: False})
-_SUPPORTED_VERSIONS = [
-    tfds.core.Version('3.0.1', (
-        'New split API (https://tensorflow.org/datasets/splits); fix colorization (all RGB) and '
-        'format (all jpeg); use TAR_STREAM.')),
-]
+_VERSION = tfds.core.Version('3.1.0')
+
+# Version history:
+# 3.1.0: Implement missing corruptions. Fix crop/resize ordering, file encoding
+# 0.0.1: Initial dataset.
 
 
 def _make_builder_configs():
   """Construct a list of BuilderConfigs.
 
-  Construct a list of 80 Imagenet2012CorruptedConfig objects, corresponding to
-  the 12 + 4 corruption types, with each type having 5 severities.
+  Construct a list of 95 Imagenet2012CorruptedConfig objects, corresponding to
+  the 15 + 4 corruption types, with each type having 5 severities.
 
   Returns:
-    A list of  Imagenet2012CorruptedConfig objects.
+    A list of 95 Imagenet2012CorruptedConfig objects.
   """
   config_list = []
   for each_corruption in BENCHMARK_CORRUPTIONS + EXTRA_CORRUPTIONS:
@@ -122,14 +129,37 @@ def _make_builder_configs():
           Imagenet2012CorruptedConfig(
               name=name_str,
               version=_VERSION,
-              supported_versions=_SUPPORTED_VERSIONS,
               description=description_str,
               corruption_type=each_corruption,
               severity=each_severity,
           ))
   return config_list
 
 
+def _decode_and_center_crop(image_bytes):
+  """Crops to center of image with padding then scales image size."""
+  shape = tf.image.extract_jpeg_shape(image_bytes)
+  image_height = shape[0]
+  image_width = shape[1]
+
+  padded_center_crop_size = tf.cast(
+      ((_IMAGE_SIZE / (_IMAGE_SIZE + _CROP_PADDING)) *
+       tf.cast(tf.minimum(image_height, image_width), tf.float32)), tf.int32)
+
+  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
+  offset_width = ((image_width - padded_center_crop_size) + 1) // 2
+  crop_window = tf.stack([
+      offset_height, offset_width, padded_center_crop_size,
+      padded_center_crop_size
+  ])
+  image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3)
+  image = tf.image.resize([image], [_IMAGE_SIZE, _IMAGE_SIZE],
+                          method=tf.image.ResizeMethod.BICUBIC)[0]
+  image = tf.cast(image, tf.int32)
+
+  return image
+
+
 class Imagenet2012Corrupted(Imagenet2012):
   """Corrupted ImageNet2012 dataset."""
   BUILDER_CONFIGS = _make_builder_configs()
@@ -145,15 +175,27 @@ def _info(self):
         builder=self,
         description=_DESCRIPTION,
         features=tfds.features.FeaturesDict({
-            'image': tfds.features.Image(),
-            'label': tfds.features.ClassLabel(names_file=names_file),
-            'file_name': tfds.features.Text(),  # Eg: 'n15075141_54.JPEG'
+            'image':
+                tfds.features.Image(
+                    shape=(_IMAGE_SIZE, _IMAGE_SIZE, 3),
+                    encoding_format='jpeg'),
+            'label':
+                tfds.features.ClassLabel(names_file=names_file),
+            'file_name':
+                tfds.features.Text(),  # Eg: 'n15075141_54.JPEG'
         }),
         supervised_keys=('image', 'label'),
         homepage='https://openreview.net/forum?id=HJz6tiCqYm',
         citation=_CITATION,
     )
 
+  def _split_generators(self, dl_manager):
+    """Filter out training split as ImageNet-C is a testing benchmark."""
+    splits = super(Imagenet2012Corrupted, self)._split_generators(dl_manager)
+
+    corruptions.FROST_FILENAMES = dl_manager.download(_FROST_FILENAMES)
+    return [s for s in splits if s.name != tfds.Split.TRAIN]
+
   def _generate_examples(self, archive, validation_labels=None):
     """Generate corrupted imagenet validation data.
 
@@ -177,7 +219,7 @@ def _generate_examples(self, archive, validation_labels=None):
     gen_fn = super(Imagenet2012Corrupted, self)._generate_examples
     for key, example in gen_fn(archive, validation_labels):
       with tf.Graph().as_default():
-        tf_img = tf.image.decode_jpeg(example['image'].read(), channels=3)
+        tf_img = _decode_and_center_crop(example['image'].read())
         image_np = tfds.as_numpy(tf_img)
       example['image'] = self._get_corrupted_example(image_np)
 
@@ -196,18 +238,22 @@ def _get_corrupted_example(self, x):
     """
     corruption_type = self.builder_config.corruption_type
     severity = self.builder_config.severity
+    x = np.clip(x, 0, 255)
 
     return {
         'gaussian_noise': corruptions.gaussian_noise,
         'shot_noise': corruptions.shot_noise,
         'impulse_noise': corruptions.impulse_noise,
         'defocus_blur': corruptions.defocus_blur,
-        'frosted_glass_blur': corruptions.frosted_glass_blur,
+        'glass_blur': corruptions.glass_blur,
+        'motion_blur': corruptions.motion_blur,
         'zoom_blur': corruptions.zoom_blur,
+        'snow': corruptions.snow,
+        'frost': corruptions.frost,
         'fog': corruptions.fog,
         'brightness': corruptions.brightness,
         'contrast': corruptions.contrast,
-        'elastic': corruptions.elastic,
+        'elastic_transform': corruptions.elastic_transform,
         'pixelate': corruptions.pixelate,
         'jpeg_compression': corruptions.jpeg_compression,
         'gaussian_blur': corruptions.gaussian_blur,