Repository: kornia/kornia

Issue:
gray_to_(rgb,bgr)
## 🚀 Feature
I have a grayscale image, and I'd like to use it with the augmentation functions which only support 3 channel images.

I'm looking for the equivalent of ``cv2.cvtColor(cv2.COLOR_GRAY2BGR, gray)``.





Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8aee6e71..7405f155 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,7 +16,7 @@ repos:
       - id: pretty-format-json
 
   - repo: https://github.com/psf/black
-    rev: 21.6b0
+    rev: 21.7b0
     hooks:
       - id: black
 
diff --git a/docs/generate_example_images.py b/docs/generate_example_images.py
index 552f184f..2af13984 100644
--- a/docs/generate_example_images.py
+++ b/docs/generate_example_images.py
@@ -128,6 +128,7 @@ def main():
 
     mod = importlib.import_module("kornia.color")
     color_transforms_list: dict = {
+        "grayscale_to_rgb": ((), 1),
         "rgb_to_bgr": ((), 1),
         "rgb_to_grayscale": ((), 1),
         "rgb_to_hsv": ((), 1),
@@ -144,7 +145,10 @@ def main():
     for fn_name, (args, num_samples) in color_transforms_list.items():
         # import function and apply
         fn = getattr(mod, fn_name)
-        out = fn(img2, *args)
+        if fn_name == "grayscale_to_rgb":
+            out = fn(K.rgb_to_grayscale(img2), *args)
+        else:
+            out = fn(img2, *args)
         # perform normalization to visualize
         if fn_name == "rgb_to_lab":
             out = out[:, :1] / 100.0
@@ -156,7 +160,12 @@ def main():
         if out.shape[1] != 3:
             out = out.repeat(1, 3, 1, 1)
         # save the output image
-        out = torch.cat([img2[0], *[out[i] for i in range(out.size(0))]], dim=-1)
+        if fn_name == "grayscale_to_rgb":
+            out = torch.cat(
+                [K.rgb_to_grayscale(img2[0]).repeat(3, 1, 1), *[out[i] for i in range(out.size(0))]], dim=-1
+            )
+        else:
+            out = torch.cat([img2[0], *[out[i] for i in range(out.size(0))]], dim=-1)
         out_np = K.utils.tensor_to_image((out * 255.0).byte())
         cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
         sig = f"{fn_name}({', '.join([str(a) for a in args])})"
diff --git a/docs/source/color.rst b/docs/source/color.rst
index 1f67b87f..880c0fc6 100644
--- a/docs/source/color.rst
+++ b/docs/source/color.rst
@@ -5,6 +5,7 @@ kornia.color
 
 The functions in this section perform various color space conversions.
 
+.. autofunction:: grayscale_to_rgb
 .. autofunction:: rgb_to_bgr
 .. autofunction:: rgb_to_grayscale
 .. autofunction:: rgb_to_hsv
@@ -42,6 +43,7 @@ The functions in this section perform various color space conversions.
 Modules
 -------
 
+.. autoclass:: GrayscaleToRgb
 .. autoclass:: RgbToGrayscale
 .. autoclass:: BgrToGrayscale
 .. autoclass:: RgbToHsv
diff --git a/kornia/__init__.py b/kornia/__init__.py
index d85320f4..d3555b9c 100644
--- a/kornia/__init__.py
+++ b/kornia/__init__.py
@@ -15,6 +15,7 @@ from kornia.color import (
     bgr_to_grayscale,
     bgr_to_rgb,
     bgr_to_rgba,
+    grayscale_to_rgb,
     hls_to_rgb,
     hsv_to_rgb,
     lab_to_rgb,
diff --git a/kornia/color/__init__.py b/kornia/color/__init__.py
index a8ec9077..6f4ed9c7 100644
--- a/kornia/color/__init__.py
+++ b/kornia/color/__init__.py
@@ -1,4 +1,4 @@
-from .gray import bgr_to_grayscale, BgrToGrayscale, rgb_to_grayscale, RgbToGrayscale
+from .gray import bgr_to_grayscale, BgrToGrayscale, grayscale_to_rgb, GrayscaleToRgb, rgb_to_grayscale, RgbToGrayscale
 from .hls import hls_to_rgb, HlsToRgb, rgb_to_hls, RgbToHls
 from .hsv import hsv_to_rgb, HsvToRgb, rgb_to_hsv, RgbToHsv
 from .lab import lab_to_rgb, LabToRgb, rgb_to_lab, RgbToLab
@@ -26,6 +26,7 @@ from .ycbcr import rgb_to_ycbcr, RgbToYcbcr, ycbcr_to_rgb, YcbcrToRgb
 from .yuv import rgb_to_yuv, RgbToYuv, yuv_to_rgb, YuvToRgb
 
 __all__ = [
+    "grayscale_to_rgb",
     "rgb_to_grayscale",
     "bgr_to_grayscale",
     "bgr_to_rgb",
diff --git a/kornia/color/gray.py b/kornia/color/gray.py
index 10bcb810..671f0eb3 100644
--- a/kornia/color/gray.py
+++ b/kornia/color/gray.py
@@ -6,6 +6,33 @@ import torch.nn as nn
 from kornia.color.rgb import bgr_to_rgb
 
 
+def grayscale_to_rgb(image: torch.Tensor) -> torch.Tensor:
+    r"""Convert a grayscale image to RGB version of image.
+
+    .. image:: _static/img/grayscale_to_rgb.png
+
+    The image data is assumed to be in the range of (0, 1).
+
+    Args:
+        image: grayscale image to be converted to RGB with shape :math:`(*,1,H,W)`.
+    Returns:
+        RGB version of the image with shape :math:`(*,3,H,W)`.
+
+    Example:
+        >>> input = torch.randn(2, 1, 4, 5)
+        >>> gray = grayscale_to_rgb(input) # 2x3x4x5
+    """
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. " f"Got {type(image)}")
+    if image.dim() < 3 or image.size(-3) != 1:
+        raise ValueError(f"Input size must have a shape of (*, 1, H, W). " f"Got {image.shape}.")
+    rgb: torch.Tensor = torch.cat([image, image, image], dim=-3)
+    image_is_float: bool = torch.is_floating_point(image)
+    if not image_is_float:
+        warnings.warn("Input image is not of float dtype. Got {}".format(image.dtype))
+    return rgb
+
+
 def rgb_to_grayscale(
     image: torch.Tensor, rgb_weights: torch.Tensor = torch.tensor([0.299, 0.587, 0.114])
 ) -> torch.Tensor:
@@ -83,6 +110,31 @@ def bgr_to_grayscale(image: torch.Tensor) -> torch.Tensor:
     return gray
 
 
+class GrayscaleToRgb(nn.Module):
+    r"""Module to convert a grayscale image to RGB version of image.
+
+    The image data is assumed to be in the range of (0, 1).
+
+    Shape:
+        - image: :math:`(*, 1, H, W)`
+        - output: :math:`(*, 3, H, W)`
+
+    reference:
+        https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html
+
+    Example:
+        >>> input = torch.rand(2, 1, 4, 5)
+        >>> rgb = GrayscaleToRgb()
+        >>> output = rgb(input)  # 2x3x4x5
+    """
+
+    def __init__(self) -> None:
+        super(GrayscaleToRgb, self).__init__()
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:  # type: ignore
+        return grayscale_to_rgb(image)
+
+
 class RgbToGrayscale(nn.Module):
     r"""Module to convert a RGB image to grayscale version of image.
 
diff --git a/kornia/jit/__init__.py b/kornia/jit/__init__.py
index 50ed3a81..52e99273 100644
--- a/kornia/jit/__init__.py
+++ b/kornia/jit/__init__.py
@@ -4,6 +4,7 @@ import kornia as K
 
 # expose functions to torch.jit
 # TODO: find an automatic way to do this
+grayscale_to_rgb = torch.jit.script(K.color.grayscale_to_rgb)
 rgb_to_grayscale = torch.jit.script(K.color.rgb_to_grayscale)
 bgr_to_grayscale = torch.jit.script(K.color.bgr_to_grayscale)
 


Test patch:
diff --git a/test/color/test_gray.py b/test/color/test_gray.py
index 2c2347d6..a4eef7ee 100644
--- a/test/color/test_gray.py
+++ b/test/color/test_gray.py
@@ -7,6 +7,103 @@ from kornia.testing import BaseTester  # test utils
 from kornia.testing import assert_close
 
 
+class TestGrayscaleToRgb(BaseTester):
+    def test_smoke(self, device, dtype):
+        C, H, W = 1, 4, 5
+        img = torch.rand(C, H, W, device=device, dtype=dtype)
+        assert isinstance(kornia.grayscale_to_rgb(img), torch.Tensor)
+
+    @pytest.mark.parametrize("batch_size, height, width", [(1, 3, 4), (2, 2, 4), (3, 4, 1)])
+    def test_cardinality(self, device, dtype, batch_size, height, width):
+        img = torch.ones(batch_size, 1, height, width, device=device, dtype=dtype)
+        assert kornia.grayscale_to_rgb(img).shape == (batch_size, 3, height, width)
+
+    def test_exception(self, device, dtype):
+        with pytest.raises(TypeError):
+            assert kornia.color.grayscale_to_rgb([0.0])
+
+        with pytest.raises(ValueError):
+            img = torch.ones(1, 1, device=device, dtype=dtype)
+            assert kornia.color.grayscale_to_rgb(img)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(2, 1, 1, device=device, dtype=dtype)
+            assert kornia.color.grayscale_to_rgb(img)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(1, 3, 1, 1, device=device, dtype=dtype)
+            assert kornia.color.grayscale_to_rgb(img)
+
+    def test_opencv(self, device, dtype):
+        data = torch.tensor(
+            [
+                [
+                    [0.4684734, 0.8954562, 0.6064363, 0.5236061, 0.6106016],
+                    [0.1709944, 0.5133104, 0.7915002, 0.5745703, 0.1680204],
+                    [0.5279005, 0.6092287, 0.3034387, 0.5333768, 0.6064113],
+                    [0.3503858, 0.5720159, 0.7052018, 0.4558409, 0.3261529],
+                    [0.6988886, 0.5897652, 0.6532392, 0.7234108, 0.7218805],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+        # Output data generated with OpenCV 4.5.2: cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB)
+        expected = torch.tensor(
+            [
+                [
+                    [0.4684734, 0.8954562, 0.6064363, 0.5236061, 0.6106016],
+                    [0.1709944, 0.5133104, 0.7915002, 0.5745703, 0.1680204],
+                    [0.5279005, 0.6092287, 0.3034387, 0.5333768, 0.6064113],
+                    [0.3503858, 0.5720159, 0.7052018, 0.4558409, 0.3261529],
+                    [0.6988886, 0.5897652, 0.6532392, 0.7234108, 0.7218805],
+                ],
+                [
+                    [0.4684734, 0.8954562, 0.6064363, 0.5236061, 0.6106016],
+                    [0.1709944, 0.5133104, 0.7915002, 0.5745703, 0.1680204],
+                    [0.5279005, 0.6092287, 0.3034387, 0.5333768, 0.6064113],
+                    [0.3503858, 0.5720159, 0.7052018, 0.4558409, 0.3261529],
+                    [0.6988886, 0.5897652, 0.6532392, 0.7234108, 0.7218805],
+                ],
+                [
+                    [0.4684734, 0.8954562, 0.6064363, 0.5236061, 0.6106016],
+                    [0.1709944, 0.5133104, 0.7915002, 0.5745703, 0.1680204],
+                    [0.5279005, 0.6092287, 0.3034387, 0.5333768, 0.6064113],
+                    [0.3503858, 0.5720159, 0.7052018, 0.4558409, 0.3261529],
+                    [0.6988886, 0.5897652, 0.6532392, 0.7234108, 0.7218805],
+                ],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+        img_rgb = kornia.grayscale_to_rgb(data)
+        assert_close(img_rgb, expected)
+
+    @pytest.mark.grad
+    def test_gradcheck(self, device, dtype):
+        B, C, H, W = 2, 1, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=torch.float64, requires_grad=True)
+        assert gradcheck(kornia.grayscale_to_rgb, (img,), raise_exception=True)
+
+    @pytest.mark.jit
+    def test_jit(self, device, dtype):
+        B, C, H, W = 2, 1, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=dtype)
+        op = kornia.grayscale_to_rgb
+        op_jit = kornia.jit.grayscale_to_rgb
+        assert_close(op(img), op_jit(img))
+
+    @pytest.mark.nn
+    def test_module(self, device, dtype):
+        B, C, H, W = 2, 1, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=dtype)
+        gray_ops = kornia.color.GrayscaleToRgb().to(device, dtype)
+        gray_fcn = kornia.color.grayscale_to_rgb
+        assert_close(gray_ops(img), gray_fcn(img))
+
+
 class TestRgbToGrayscale(BaseTester):
     def test_smoke(self, device, dtype):
         C, H, W = 3, 4, 5