Repository: kornia/kornia

Issue:
[Feat] Raw to RGB
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Hi. Thank you for the great work! Is there any plan or possible solution for raw-to-RGB (Bayer Pattern to RGB) transformation? The corresponding OpenCV function is as follows:
`img_bgr = cv2.cvtColor(img_bayer_bg.astype(np.uint8), cv2.COLOR_BAYER_BG2BGR)`
Thank you!

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
OpenCV raw-to-RGB function in a differentiable way.

## Pitch

<!-- A clear and concise description of what you want to happen. -->

A way to transform raw image to RGB image in a differentiable manner.

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
@ycjing this functionality it's definitely inside the scope of the library. We can add it to the roadmap for the upcoming releases
Hi @edgarriba 

Thank you for the response! I appreciate it.

Best,
Yongcheng
@cceyda is this what you did with `rgb_to_linear` ?
@edgarriba  No, rgb_to_linear reverses the [gamma correction](https://stackoverflow.com/questions/12524623/what-are-the-practical-differences-when-working-with-colors-in-a-linear-vs-a-no) on an sRGB image.
Unfortunately I'm not familiar with Bayer Patterns, which looks to be a low level processing thing done by digital cameras 🤷 
I'm curious about your use case though @ycjing

Resolving patch:
diff --git a/kornia/color/__init__.py b/kornia/color/__init__.py
--- a/kornia/color/__init__.py
+++ b/kornia/color/__init__.py
@@ -3,6 +3,7 @@
 from .hsv import hsv_to_rgb, HsvToRgb, rgb_to_hsv, RgbToHsv
 from .lab import lab_to_rgb, LabToRgb, rgb_to_lab, RgbToLab
 from .luv import luv_to_rgb, LuvToRgb, rgb_to_luv, RgbToLuv
+from .raw import CFA, raw_to_rgb, RawToRgb, rgb_to_raw, RgbToRaw
 from .rgb import (
     bgr_to_rgb,
     bgr_to_rgba,
@@ -84,4 +85,9 @@
     "LuvToRgb",
     "LabToRgb",
     "RgbToLab",
+    "RgbToRaw",
+    "RawToRgb",
+    "raw_to_rgb",
+    "rgb_to_raw",
+    "CFA",
 ]
diff --git a/kornia/color/raw.py b/kornia/color/raw.py
new file mode 100644
--- /dev/null
+++ b/kornia/color/raw.py
@@ -0,0 +1,257 @@
+from enum import Enum
+
+import torch
+import torch.nn as nn
+
+
+class CFA(Enum):
+    r"""Defines the configuration of the color filter array.
+
+    So far only bayer images is supported and the enum sets the pixel order for bayer. Note that this can change due
+    to things like rotations and cropping of images. Take care if including the translations in pipeline.
+    This implementations is optimized to be reasonably fast, look better than simple nearest neighbour.
+    On top of this care is taken to make it reversible going raw -> rgb -> raw. the raw samples remain intact
+    during conversion and only unknown samples are interpolated.
+
+    The names are based on the OpenCV convention where the BG indicates pixel 1,1 (counting from 0,0) is
+    blue and its neighbour to the right is green. In that case the top left pixel is red. Other options are GB, RG and
+    GR
+
+    reference:
+        https://en.wikipedia.org/wiki/Color_filter_array
+    """
+
+    BG = 0
+    GB = 1
+    RG = 2
+    GR = 3
+
+
+def raw_to_rgb(image: torch.Tensor, cfa: CFA) -> torch.Tensor:
+    r"""Convert a raw bayer image to RGB version of image.
+
+    We are assuming a CFA with 2 green, 1 red, 1 blue. A bilinear interpolation is used for R/G and a fix convolution
+    for the green pixels. To simplify calculations we expect the Height Width to be evenly divisible by 2.
+
+    The image data is assumed to be in the range of (0, 1). Image H/W is assumed to be evenly divisible by 2.
+    for simplicity reasons
+
+    Args:
+        image: raw image to be converted to RGB with shape :math:`(*,1,H,W)`.
+        cfa: The configuration of the color filter.
+    Returns:
+        RGB version of the image with shape :math:`(*,3,H,W)`.
+
+    Example:
+        >>> rawinput = torch.randn(2, 1, 4, 6)
+        >>> rgb = raw_to_rgb(rawinput, CFA.RG) # 2x3x4x6
+    """
+    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}.")
+
+    if len(image.shape) < 2 or image.shape[-2] % 2 == 1 or image.shape[-1] % 2 == 1:
+        raise ValueError(f"Input H&W must be evenly disible by 2. Got {image.shape}")
+
+    imagesize = image.size()
+
+    image = image.view(-1, 1, image.shape[-2], image.shape[-1])
+
+    # BG is defined as pel 1,1 being blue, that is the top left is actually green. This matches
+    # opencv naming so makes sense to keep
+    if cfa == CFA.BG:
+        r = image[..., :, ::2, ::2]
+        b = image[..., :, 1::2, 1::2]
+        rpad = (0, 1, 0, 1)
+        bpad = (1, 0, 1, 0)
+    elif cfa == CFA.GB:
+        r = image[..., :, ::2, 1::2]
+        b = image[..., :, 1::2, ::2]
+        rpad = (1, 0, 0, 1)
+        bpad = (0, 1, 1, 0)
+    elif cfa == CFA.RG:
+        r = image[..., :, 1::2, 1::2]
+        b = image[..., :, ::2, ::2]
+        rpad = (1, 0, 1, 0)
+        bpad = (0, 1, 0, 1)
+    elif cfa == CFA.GR:
+        r = image[..., :, 1::2, ::2]
+        b = image[..., :, ::2, 1::2]
+        rpad = (0, 1, 1, 0)
+        bpad = (1, 0, 0, 1)
+    else:
+        raise ValueError(f"Unsupported CFA " f"Got {cfa}.")
+
+    # upscaling r and b with bi-linear gives reasonable quality
+    # Note that depending on where these are sampled we need to pad appropriately
+    # the bilinear filter will pretty much be based on for example this layout (RG)
+    # (which needs to be padded bottom right)
+    # +-+-+
+    # |B| |
+    # | | |
+    # +-+-+
+    # While in this layout we need to pad with additional B samples top left to
+    # make sure we interpolate from the correct position
+    # +-+-+
+    # | | |
+    # | |B|
+    # +-+-+
+    # For an image like this (3x2 blue pixels)
+    # +------+
+    # |B B B |
+    # |      |
+    # |B B B |
+    # |      |
+    # +------+
+    # It needs to be expanded to this (4x3 pixels scaled to 7x5 for correct interpolation)
+    # +-------+
+    # |B B B b|
+    # |       |
+    # |B B B b|
+    # |       |
+    # |b b b b|
+    # +-------+
+    # and we crop the area afterwards. This is since the interpolation will be between first and last pixel
+    # evenly spaced between them while the B/R samples will be missing in the corners were they are assumed to exist
+    # Further we need to do align_corners to start the interpolation from the middle of the samples in the corners, that
+    # way we get to keep the known blue samples across the whole image
+    rpadded = torch.nn.functional.pad(r, rpad, 'replicate')
+    bpadded = torch.nn.functional.pad(b, bpad, 'replicate')
+    # use explicit padding instead of conv2d padding to be able to use reflect which mirror the correct colors
+    # for a 2x2 bayer filter
+    gpadded = torch.nn.functional.pad(image, (1, 1, 1, 1), 'reflect')
+
+    ru = torch.nn.functional.interpolate(rpadded, size=(image.shape[-2] + 1, image.shape[-1] + 1),
+                                         mode='bilinear', align_corners=True)
+    bu = torch.nn.functional.interpolate(bpadded, size=(image.shape[-2] + 1, image.shape[-1] + 1),
+                                         mode='bilinear', align_corners=True)
+
+    # remove the extra padding
+    ru = torch.nn.functional.pad(ru, [-x for x in rpad])
+    bu = torch.nn.functional.pad(bu, [-x for x in bpad])
+
+    # all unknown pixels are the average of the nearby green samples
+    kernel = torch.tensor([[[[0.0, 0.25, 0.0], [0.25, 0.0, 0.25], [0.0, 0.25, 0.0]]]],
+                          dtype=image.dtype, device=image.device)
+
+    # This is done on all samples but result for the known green samples is then overwritten by the input
+    gu = torch.nn.functional.conv2d(gpadded, kernel)
+
+    # overwrite the already known samples which otherwise have values from r/b
+    # this depends on the CFA configuration
+    if cfa == CFA.BG:
+        gu[:, :, ::2, 1::2] = image[:, :, ::2, 1::2]
+        gu[:, :, 1::2, ::2] = image[:, :, 1::2, ::2]
+    elif cfa == CFA.GB:
+        gu[:, :, ::2, ::2] = image[:, :, ::2, ::2]
+        gu[:, :, 1::2, 1::2] = image[:, :, 1::2, 1::2]
+    elif cfa == CFA.RG:
+        gu[:, :, 1::2, ::2] = image[:, :, 1::2, ::2]
+        gu[:, :, ::2, 1::2] = image[:, :, ::2, 1::2]
+    elif cfa == CFA.GR:
+        gu[:, :, 1::2, 1::2] = image[:, :, 1::2, 1::2]
+        gu[:, :, ::2, ::2] = image[:, :, ::2, ::2]
+    else:
+        raise ValueError(f"Unsupported CFA " f"Got {cfa}.")
+
+    ru = ru.view(imagesize)
+    gu = gu.view(imagesize)
+    bu = bu.view(imagesize)
+
+    rgb: torch.Tensor = torch.cat([ru, gu, bu], dim=-3)
+
+    return rgb
+
+
+def rgb_to_raw(image: torch.Tensor, cfa: CFA) -> torch.Tensor:
+    r"""Convert a RGB image to RAW version of image with the specified color filter array.
+
+    The image data is assumed to be in the range of (0, 1).
+
+    Args:
+        image: RGB image to be converted to bayer raw with shape :math:`(*,3,H,W)`.
+        cfa: Which color filter array do we want the output to mimic. I.e. which pixels are red/green/blue.
+
+    Returns:
+        raw version of the image with shape :math:`(*,1,H,W)`.
+
+    Example:
+        >>> rgbinput = torch.rand(2, 3, 4, 6)
+        >>> raw = rgb_to_raw(rgbinput, CFA.BG) # 2x1x4x6
+    """
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(image)}")
+
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    # pick the tensor with green pixels
+    # clone to make sure grad works
+    output: torch.Tensor = image[..., 1:2, :, :].clone()
+
+    # overwrite the r/b positions (depending on the cfa configuration) with blue/red pixels
+    if cfa == CFA.BG:
+        output[..., :, ::2, ::2] = image[..., 0:1, ::2, ::2]  # red
+        output[..., :, 1::2, 1::2] = image[..., 2:3, 1::2, 1::2]  # blue
+    elif cfa == CFA.GB:
+        output[..., :, ::2, 1::2] = image[..., 0:1, ::2, 1::2]  # red
+        output[..., :, 1::2, ::2] = image[..., 2:3, 1::2, ::2]  # blue
+    elif cfa == CFA.RG:
+        output[..., :, 1::2, 1::2] = image[..., 0:1, 1::2, 1::2]  # red
+        output[..., :, ::2, ::2] = image[..., 2:3, ::2, ::2]  # blue
+    elif cfa == CFA.GR:
+        output[..., :, 1::2, ::2] = image[..., 0:1, 1::2, ::2]  # red
+        output[..., :, ::2, 1::2] = image[..., 2:3, ::2, 1::2]  # blue
+
+    return output
+
+
+class RawToRgb(nn.Module):
+    r"""Module to convert a bayer raw 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)`
+
+    Example:
+        >>> rawinput = torch.rand(2, 1, 4, 6)
+        >>> rgb = RawToRgb(CFA.RG)
+        >>> output = rgb(rawinput)  # 2x3x4x5
+    """
+
+    def __init__(self, cfa: CFA) -> None:
+        super().__init__()
+        self.cfa = cfa
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        return raw_to_rgb(image, cfa=self.cfa)
+
+
+class RgbToRaw(nn.Module):
+    r"""Module to convert a RGB image to bayer raw version of image.
+
+    The image data is assumed to be in the range of (0, 1).
+
+    Shape:
+        - image: :math:`(*, 3, H, W)`
+        - output: :math:`(*, 1, H, W)`
+
+    reference:
+        https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html
+
+    Example:
+        >>> rgbinput = torch.rand(2, 3, 4, 6)
+        >>> raw = RgbToRaw(CFA.GB)
+        >>> output = raw(rgbinput)  # 2x1x4x6
+    """
+
+    def __init__(self, cfa: CFA) -> None:
+        super().__init__()
+        self.cfa = cfa
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        return rgb_to_raw(image, cfa=self.cfa)


Test patch:
diff --git a/test/color/test_raw.py b/test/color/test_raw.py
new file mode 100644
--- /dev/null
+++ b/test/color/test_raw.py
@@ -0,0 +1,195 @@
+import warnings
+
+import pytest
+import torch
+from torch.autograd import gradcheck
+
+import kornia
+from kornia.testing import BaseTester  # test utils
+from kornia.testing import assert_close
+from packaging import version
+
+
+class TestRawToRgb(BaseTester):
+    def test_smoke(self, device, dtype):
+        C, H, W = 1, 4, 6
+        img = torch.rand(C, H, W, device=device, dtype=dtype)
+        assert isinstance(kornia.color.raw_to_rgb(img, kornia.color.CFA.BG), torch.Tensor)
+
+    @pytest.mark.parametrize("batch_size, height, width", [(1, 6, 4), (2, 2, 4), (3, 4, 2)])
+    def test_cardinality(self, device, dtype, batch_size, height, width):
+        img = torch.ones(batch_size, 1, height, width, device=device, dtype=dtype)
+        assert kornia.color.raw_to_rgb(img, kornia.color.CFA.BG).shape == (batch_size, 3, height, width)
+
+    def test_exception(self, device, dtype):
+        with pytest.raises(TypeError):
+            assert kornia.color.raw_to_rgb([0.0], kornia.color.CFA.BG)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(1, 1, device=device, dtype=dtype)
+            assert kornia.color.raw_to_rgb(img, kornia.color.CFA.GB)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(2, 1, 1, device=device, dtype=dtype)
+            assert kornia.color.raw_to_rgb(img, kornia.color.CFA.RG)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(1, 3, 1, 1, device=device, dtype=dtype)
+            assert kornia.color.raw_to_rgb(img, kornia.color.CFA.GR)
+
+        # dimensionality test
+        with pytest.raises(ValueError):
+            img = torch.ones(3, 2, 1, device=device, dtype=dtype)
+            assert kornia.color.raw_to_rgb(img, kornia.color.CFA.GR)
+
+        # dimensionality test
+        with pytest.raises(ValueError):
+            img = torch.ones(3, 1, 2, device=device, dtype=dtype)
+            assert kornia.color.raw_to_rgb(img, kornia.color.CFA.GR)
+
+    # With he current implementations we should get back an identical raw representation when doing raw -> rgb -> raw
+    # Note that with more advanced implementations this may not necessarily be true or desirable
+    def test_forth_and_back(self, device, dtype):  # skipcq: PYL-R0201
+        data = torch.rand(1, 80, 80, device=device, dtype=dtype)
+        raw = kornia.color.rgb_to_raw
+        rgb = kornia.color.raw_to_rgb
+
+        for x in kornia.color.CFA:
+            data_out = raw(rgb(data, cfa=x), cfa=x)
+            assert_close(data_out, data)
+
+    # make sure different cfas are actually different
+    def test_cfas_not_the_same(self, device, dtype):  # skipcq: PYL-R0201
+        data = torch.rand(1, 16, 16, device=device, dtype=dtype)
+        assert torch.max(kornia.color.raw_to_rgb(data, kornia.color.CFA.BG)
+                         - kornia.color.raw_to_rgb(data, kornia.color.CFA.RG)) > 0.0
+
+    # The outcome will be very different for different implementations
+    # Here we compare against a current baseline, it is safe to update this if the underlying algorithm changes
+    def test_functional(self, device, dtype):  # skipcq: PYL-R0201
+        data = torch.tensor(
+            [
+                [
+                    [1, 0.5, 0.2, 0.4],
+                    [0.75, 0.25, 0.8, 0.3],
+                    [0.65, 0.15, 0.7, 0.2],
+                    [0.55, 0.5, 0.6, 0.1],
+                ]
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        # checked by hand as correct interpolation. Note the ugly replication that happens for Red on the last column
+        # and row. We shall accept to live with that
+        expected = torch.tensor([[[1.0000, 0.6000, 0.2000, 0.2000],
+                                  [0.8250, 0.6375, 0.4500, 0.4500],
+                                  [0.6500, 0.6750, 0.7000, 0.7000],
+                                  [0.6500, 0.6750, 0.7000, 0.7000]],
+                                 [[0.6250, 0.5000, 0.6250, 0.4000],
+                                  [0.7500, 0.5500, 0.8000, 0.5500],
+                                  [0.4000, 0.1500, 0.4375, 0.2000],
+                                  [0.5500, 0.3625, 0.6000, 0.4000]],
+                                 [[0.2500, 0.2500, 0.2750, 0.3000],
+                                  [0.2500, 0.2500, 0.2750, 0.3000],
+                                  [0.3750, 0.3750, 0.2875, 0.2000],
+                                  [0.5000, 0.5000, 0.3000, 0.1000]]],
+                                device=device,
+                                dtype=dtype,
+                                )
+
+        img_rgb = kornia.color.raw_to_rgb(data, kornia.color.raw.CFA.BG)
+        assert_close(img_rgb, expected)
+
+    # If we roll the data and the different CFAs they give the same result (expect on edges!)
+    def test_cfa_on_rolled(self, device, dtype):  # skipcq: PYL-R0201
+        data = torch.rand(1, 1, 8, 8, device=device, dtype=dtype)
+        bgres = kornia.color.raw_to_rgb(data, kornia.color.raw.CFA.BG)
+        gbres = kornia.color.raw_to_rgb(data.roll((0, 1), (-2, -1)), kornia.color.raw.CFA.GB)
+        grres = kornia.color.raw_to_rgb(data.roll((1, 0), (-2, -1)), kornia.color.raw.CFA.GR)
+        rgres = kornia.color.raw_to_rgb(data.roll((1, 1), (-2, -1)), kornia.color.raw.CFA.RG)
+
+        assert_close(bgres[:, :, 1:5, 1:5], gbres[:, :, 1:5, 2:6])
+        assert_close(bgres[:, :, 1:5, 1:5], grres[:, :, 2:6, 1:5])
+        assert_close(bgres[:, :, 1:5, 1:5], rgres[:, :, 2:6, 2:6])
+
+    @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.color.raw_to_rgb, (img, kornia.color.raw.CFA.BG), raise_exception=True)
+
+    @pytest.mark.jit
+    def test_jit(self, device, dtype):
+        if version.parse(torch.__version__) < version.parse('1.7.0'):
+            warnings.warn(
+                "This test is not compatible with pytorch < 1.7.0. This message will be removed as soon as we do not "
+                "support pytorch 1.6.0. `rgb_to_hls()` method for pytorch < 1.7.0 version cannot be compiled with JIT.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return
+        B, C, H, W = 2, 1, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=dtype)
+        op = kornia.color.raw_to_rgb
+        op_jit = torch.jit.script(op)
+        assert_close(op(img, kornia.color.raw.CFA.BG), op_jit(img, kornia.color.raw.CFA.BG))
+
+    @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)
+        raw_ops = kornia.color.RawToRgb(kornia.color.raw.CFA.BG).to(device, dtype)
+        raw_fcn = kornia.color.raw_to_rgb
+        assert_close(raw_ops(img), raw_fcn(img, kornia.color.raw.CFA.BG))
+
+
+class TestRgbToRaw(BaseTester):
+    def test_smoke(self, device, dtype):
+        C, H, W = 3, 4, 6
+        img = torch.rand(C, H, W, device=device, dtype=dtype)
+        assert isinstance(kornia.color.rgb_to_raw(img, kornia.color.raw.CFA.BG), 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, 3, height, width, device=device, dtype=dtype)
+        assert kornia.color.rgb_to_raw(img, kornia.color.raw.CFA.GR).shape == (batch_size, 1, height, width)
+
+    def test_exception(self, device, dtype):
+        with pytest.raises(TypeError):
+            assert kornia.color.rgb_to_raw([0.0], kornia.color.raw.CFA.RG)
+
+        with pytest.raises(ValueError):
+            img = torch.ones(1, 1, device=device, dtype=dtype)
+            assert kornia.color.rgb_to_raw(img, kornia.color.raw.CFA.BG)
+
+        # Reverse test in rawtorgb is sufficient functional test
+
+    @pytest.mark.grad
+    def test_gradcheck(self, device, dtype):
+        B, C, H, W = 2, 3, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=torch.float64, requires_grad=True)
+        assert gradcheck(kornia.color.rgb_to_raw, (img, kornia.color.raw.CFA.BG), raise_exception=True)
+
+    @pytest.mark.jit
+    def test_jit(self, device, dtype):
+        if version.parse(torch.__version__) < version.parse('1.7.0'):
+            warnings.warn(
+                "This test is not compatible with pytorch < 1.7.0. This message will be removed as soon as we do not "
+                "support pytorch 1.6.0. `rgb_to_hls()` method for pytorch < 1.7.0 version cannot be compiled with JIT.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return
+        B, C, H, W = 2, 3, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=dtype)
+        op = kornia.color.rgb_to_raw
+        op_jit = torch.jit.script(op)
+        assert_close(op(img, kornia.color.raw.CFA.BG), op_jit(img, kornia.color.raw.CFA.BG))
+
+    @pytest.mark.nn
+    def test_module(self, device, dtype):
+        B, C, H, W = 2, 3, 4, 4
+        img = torch.ones(B, C, H, W, device=device, dtype=dtype)
+        raw_ops = kornia.color.RgbToRaw(kornia.color.raw.CFA.BG).to(device, dtype)
+        raw_fcn = kornia.color.rgb_to_raw
+        assert_close(raw_ops(img), raw_fcn(img, kornia.color.raw.CFA.BG))