updated docstring, fixed motion blur

imagecorruptions/__init__.py

@@ -12,31 +12,61 @@
                    corruption_tuple}
 
 
-def corrupt(x, severity=1, corruption_name=None, corruption_number=-1):
-    """
-    :param x: image to corrupt; a 224x224x3 numpy array in [0, 255]
-    :param severity: strength with which to corrupt x; an integer in [0, 5]
-    :param corruption_name: specifies which corruption function to call;
-    must be one of 'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
-                    'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
-                    'brightness', 'contrast', 'elastic_transform', 'pixelate', 'jpeg_compression',
-                    'speckle_noise', 'gaussian_blur', 'spatter', 'saturate';
-                    the last four are validation functions
-    :param corruption_number: the position of the corruption_name in the above list;
-    an integer in [0, 18]; useful for easy looping; 15, 16, 17, 18 are validation corruption numbers
-    :return: the image x corrupted by a corruption function at the given severity; same shape as input
+def corrupt(image, severity=1, corruption_name=None, corruption_number=-1):
+    """This function returns a corrupted version of the given image.
+
+    Args:
+        image (numpy.ndarray):      image to corrupt; a numpy array in [0, 255], expected datatype is np.uint8
+                                    expected shape is either (height x width x channels) or (height x width); 
+                                    width and height must be at least 32 pixels;
+                                    channels must be 1 or 3;
+        severity (int):             strength with which to corrupt the image; an integer in [1, 5]
+        corruption_name (str):      specifies which corruption function to call, must be one of
+                                        'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
+                                        'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
+                                        'brightness', 'contrast', 'elastic_transform', 'pixelate', 'jpeg_compression',
+                                        'speckle_noise', 'gaussian_blur', 'spatter', 'saturate';
+                                        the last four are validation corruptions
+        corruption_number (int):    the position of the corruption_name in the above list; an integer in [0, 18]; 
+                                        useful for easy looping; 15, 16, 17, 18 are validation corruption numbers
+    Returns:
+        numpy.ndarray:              the image corrupted by a corruption function at the given severity; same shape as input
     """
 
-    if corruption_name:
-        x_corrupted = corruption_dict[corruption_name](Image.fromarray(x),
+    if not isinstance(image, np.ndarray):
+        raise AttributeError('Expecting type(image) to be numpy.ndarray')
+    if not (image.dtype.type is np.uint8):
+        raise AttributeError('Expecting image.dtype.type to be numpy.uint8')
+
+    if not (image.ndim in [2,3]):
+        raise AttributeError('Expecting image.shape to be either (width x height) or (width x height x channels)')
+    if image.ndim == 2:
+        image = np.stack((image,)*3, axis=-1)
+
+    height, width, channels = image.shape
+
+    if (height < 32 or width < 32):
+        raise AttributeError('Image width and height must be at least 32 pixels')
+
+    if not (channels in [1,3]):
+        raise AttributeError('Expecting image to have either 1 or 3 channels (last dimension)')
+
+    if channels == 1:
+        image = np.stack((np.squeeze(image),)*3, axis=-1)
+
+    if not severity in [1,2,3,4,5]:
+        raise AttributeError('Severity must be an integer in [1, 5]')
+
+    if not (corruption_name is None):
+        image_corrupted = corruption_dict[corruption_name](Image.fromarray(image),
                                                        severity)
     elif corruption_number != -1:
-        x_corrupted = corruption_tuple[corruption_number](Image.fromarray(x),
+        image_corrupted = corruption_tuple[corruption_number](Image.fromarray(image),
                                                           severity)
     else:
         raise ValueError("Either corruption_name or corruption_number must be passed")
 
-    return np.uint8(x_corrupted)
+    return np.uint8(image_corrupted)
 
 def get_corruption_names(subset='common'):
     if subset == 'common':

imagecorruptions/corruptions.py

@@ -117,20 +117,20 @@ def getMotionBlurKernel(width, sigma):
 
 def shift(image, dx, dy):
     if(dx < 0):
-        shifted = np.roll(image, shift=image.shape[0]+dx, axis=0)
-        shifted[dx:,:] = shifted[dx-1:dx,:]
+        shifted = np.roll(image, shift=image.shape[1]+dx, axis=1)
+        shifted[:,dx:] = shifted[:,dx-1:dx]
     elif(dx > 0):
-        shifted = np.roll(image, shift=dx, axis=0)
-        shifted[:dx,:] = shifted[dx:dx+1,:]
+        shifted = np.roll(image, shift=dx, axis=1)
+        shifted[:,:dx] = shifted[:,dx:dx+1]
     else:
         shifted = image
 
     if(dy < 0):
-        shifted = np.roll(shifted, shift=image.shape[1]+dy, axis=1)
-        shifted[:,dy:] = shifted[:,dy-1:dy]
+        shifted = np.roll(shifted, shift=image.shape[0]+dy, axis=0)
+        shifted[dy:,:] = shifted[dy-1:dy,:]
     elif(dy > 0):
-        shifted = np.roll(shifted, shift=dy, axis=1)
-        shifted[:,:dy] = shifted[:,dy:dy+1]
+        shifted = np.roll(shifted, shift=dy, axis=0)
+        shifted[:dy,:] = shifted[dy:dy+1,:]
     return shifted
 
 def _motion_blur(x, radius, sigma, angle):
@@ -141,14 +141,15 @@ def _motion_blur(x, radius, sigma, angle):
 
     blurred = np.zeros_like(x, dtype=np.float32)
     for i in range(width):
-        dx = -math.ceil(((i*point[0]) / hypot) - 0.5)
-        dy = -math.ceil(((i*point[1]) / hypot) - 0.5)
+        dy = -math.ceil(((i*point[0]) / hypot) - 0.5)
+        dx = -math.ceil(((i*point[1]) / hypot) - 0.5)
+        if (np.abs(dy) >= x.shape[0] or np.abs(dx) >= x.shape[1]):
+            # simulated motion exceeded image borders
+            break
         shifted = shift(x, dx, dy)
         blurred = blurred + kernel[i] * shifted
-
     return blurred
 
-
 # /////////////// End Corruption Helpers ///////////////