PR review

Author: vfdev-5

torchvision/prototype/transforms/functional/_color.py

@@ -2,7 +2,7 @@
 from torchvision.prototype import features
 from torchvision.transforms import functional_pil as _FP, functional_tensor as _FT
 
-from ._meta import get_dimensions_image_tensor, get_num_channels_image_tensor, rgb_to_grayscale
+from ._meta import _rgb_to_gray, get_dimensions_image_tensor, get_num_channels_image_tensor
 
 
 def _blend(image1: torch.Tensor, image2: torch.Tensor, ratio: float) -> torch.Tensor:
@@ -52,7 +52,7 @@ def adjust_saturation_image_tensor(image: torch.Tensor, saturation_factor: float
     if c == 1:  # Match PIL behaviour
         return image
 
-    return _blend(image, rgb_to_grayscale(image), saturation_factor)
+    return _blend(image, _rgb_to_gray(image), saturation_factor)
 
 
 adjust_saturation_image_pil = _FP.adjust_saturation
@@ -79,7 +79,7 @@ def adjust_contrast_image_tensor(image: torch.Tensor, contrast_factor: float) ->
     if c not in [1, 3]:
         raise TypeError(f"Input image tensor permitted channel values are {[1, 3]}, but found {c}")
     dtype = image.dtype if torch.is_floating_point(image) else torch.float32
-    grayscale_image = _FT.rgb_to_grayscale(image) if c == 3 else image
+    grayscale_image = _rgb_to_gray(image) if c == 3 else image
     mean = torch.mean(grayscale_image.to(dtype), dim=(-3, -2, -1), keepdim=True)
     return _blend(image, mean, contrast_factor)
 

torchvision/prototype/transforms/functional/_meta.py

@@ -184,33 +184,13 @@ def _gray_to_rgb(grayscale: torch.Tensor) -> torch.Tensor:
     return grayscale.repeat(repeats)
 
 
-def rgb_to_grayscale(image: torch.Tensor, num_output_channels: int = 1) -> torch.Tensor:
-    if image.ndim < 3:
-        raise TypeError(f"Input image tensor should have at least 3 dimensions, but found {image.ndim}")
-
-    c = image.shape[-3]
-    if c not in [1, 3]:
-        raise TypeError(f"Input image tensor permitted channel values are {[1, 3]}, but found {c}")
-
-    if num_output_channels not in (1, 3):
-        raise ValueError("num_output_channels should be either 1 or 3")
-
-    if c == 3:
-        r, g, b = image.unbind(dim=-3)
-        l_img = (0.2989 * r).add_(g, alpha=0.587).add_(b, alpha=0.114)
-        l_img = l_img.to(image.dtype).unsqueeze(dim=-3)
-    else:
-        l_img = image.clone()
-
-    if num_output_channels == 3:
-        return l_img.expand(image.shape)
-
+def _rgb_to_gray(image: torch.Tensor) -> torch.Tensor:
+    r, g, b = image.unbind(dim=-3)
+    l_img = (0.2989 * r).add_(g, alpha=0.587).add_(b, alpha=0.114)
+    l_img = l_img.to(image.dtype).unsqueeze(dim=-3)
     return l_img
 
 
-_rgb_to_gray = rgb_to_grayscale
-
-
 def convert_color_space_image_tensor(
     image: torch.Tensor, old_color_space: ColorSpace, new_color_space: ColorSpace, copy: bool = True
 ) -> torch.Tensor: