add scriptable transforms: rgb_to_grayscale

Author: PyExtreme

torchvision/transforms/functional_tensor.py

@@ -49,24 +49,28 @@ def crop(img, top, left, height, width):
 
     return img[..., top:top + height, left:left + width]
 
-def to_grayscale(img, num_output_channels = 3):
+
+def rgb_to_grayscale(img, num_output_channels=3):
     """Convert the given RGB Image Tensor to Grayscale.
 
     Args
         img (Tensor): Image to be converted to Grayscale in the form [C, H, W].
         num_output_channels (int): denotes the number of channels to return after conversion
     Returns:
     Tensor: Grayscale image.
+
+    For RGB to Grayscale conversion, ITU-R 601-2 luma transform is performed which is
+    L = R * 0.2989 + G * 0.5870 + B * 0.1140
     """
     if not F._is_tensor_image(img):
         raise TypeError('tensor is not a torch image.')
-
+    if img.size()[0] != 3:
+        raise TypeError('Input Image doesn\'t have 3 Channels')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    img = img.to(device)
+    weights = torch.tensor([0.2989, 0.5870, 0.1140])
+    res = torch.tensordot(img, weights[:, None, None], [[0], [0]]).squeeze()
+    if num_output_channels == 1:
+        return res
     else:
-        hwc_img = img.transpose(1, 2).transpose(0, 2)
-        weights = torch.tensor([0.2989, 0.5870, 0.1140])
-        res = hwc_img[:,:,:3] * weights[None, :]
-        gray_img = res[:, :, 0] + res[:, :, 1] + res[:, :, 2]
-        if num_output_channels == 1:
-            return gray_img.int()
-        else:
-            return torch.cat((gray_img.int(), gray_img.int(), gray_img.int())).transpose(1, 2)
+        return res.repeat(3, 1, 1)