Added RandomGrayscale transform

Author: sourabhd

torchvision/transforms.py

@@ -494,6 +494,26 @@ def adjust_gamma(img, gamma, gain=1):
     img = Image.fromarray(np_img, 'RGB').convert(input_mode)
     return img
 
+def to_grayscale(img):
+    """Convert image to grayscale, repeated over three channels.
+
+    Args:
+        img (PIL Image): Image to be converted to grayscale. 
+
+    Returns:
+        PIL Image:  Grayscale version of the image, repeated over three channels.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    input_mode = img.mode
+    img = img.convert('L')
+
+    np_img = np.array(img, dtype=np.uint8)
+    np_img = np.dstack([np_img] * 3)
+
+    img = Image.fromarray(np_img, 'RGB').convert(input_mode)
+    return img
 
 class Compose(object):
     """Composes several transforms together.
@@ -1026,3 +1046,26 @@ def __call__(self, img):
         transform = self.get_params(self.brightness, self.contrast,
                                     self.saturation, self.hue)
         return transform(img)
+
+
+class RandomGrayscale(object):
+    """Randomly convert image to grayscale with a probability of p (default 0.1).
+    Args:
+        p (float): probability that image should be converted to grayscale.
+
+    """
+
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be converted to grayscale.
+
+        Returns:
+            PIL Image: Randomly grayscaled image.
+        """
+        if random.random() < self.p:
+            return to_grayscale(img)
+        return img