Repository: kornia/kornia

Issue:
[Feature request] Debayer and decimation in single step
## 🚀 Feature

A common image signal processing (ISP) step in most camera pipelines are to use a bayer filter to get color images from raw pixel values. Kornia already has conversions from bayer to color images. In computer vision it is often necessary to downscale images for computational reasons or because it yields improved results (e.g. in stereo pipelines it can be considered a regularization technique). The straightforward ISP is then to do `Bayer Image -> Color Image -> Resized Color Image`.

This ISP pipeline, however, has redundant computational steps and erodes the raw signal by successive interpolation steps.

A better approach is to do debayering and resizing/decimation in one step.

Consider the standard `rggb` pattern:

-------
R   | G1
G2 | B
-------

A fast and simple approach to get downsampled color images is to simply extract the three color channels directly to a downscaled color image:

-----
R (G1+G2)/2 B
-----

This can be very effeciently implemented in pytorch and as an additional benefit, this contains the full raw signal in the red and blue channels and only slightly erodes it in the green channel.


## Motivation

To enable more research into how ISPs affect computer vision, Kornia can implement simple and easy to use functionality for common ISP steps.

## Pitch

In the `kornia.color` module implement one function:

```
def rgb_to_raw_2x2_downscaled(image: torch.Tensor, cfa: CFA) -> torch.Tensor
```
That will return the downscaled color image from the raw bayer image.

## Alternatives

N/A

## Additional context

See following papers for more info:

* https://ieeexplore.ieee.org/abstract/document/9385089
* https://hal.archives-ouvertes.fr/hal-00683233/document

I can contribute this functionality or supervise anyone who wants to contribute.



Discussion:
@copaah check https://github.com/kornia/kornia/blob/master/kornia/color/raw.py /cc @oskarflordal @cceyda 
Yes, it should be straight forward and fairly common thing to do, I suppose. The function is essentially taking the BG generated here
https://github.com/kornia/kornia/blob/master/kornia/color/raw.py#L64
and adding the two green samples from here https://github.com/kornia/kornia/blob/master/kornia/color/raw.py#L144
and stacking the 3 channels together.
@oskarflordal 

Yes, exactly. This is also what the ROS Image Pipeline does [here](https://github.com/ros-perception/image_pipeline/blob/71d537a50bf4e7769af513f4a3d3df0d188cfa05/image_proc/src/nodelets/crop_decimate.cpp#L120).

Should be a very easy PR.

Resolving patch:
diff --git a/kornia/color/__init__.py b/kornia/color/__init__.py
--- a/kornia/color/__init__.py
+++ b/kornia/color/__init__.py
@@ -4,7 +4,7 @@
 from .hsv import HsvToRgb, RgbToHsv, hsv_to_rgb, rgb_to_hsv
 from .lab import LabToRgb, RgbToLab, lab_to_rgb, rgb_to_lab
 from .luv import LuvToRgb, RgbToLuv, luv_to_rgb, rgb_to_luv
-from .raw import CFA, RawToRgb, RgbToRaw, raw_to_rgb, rgb_to_raw
+from .raw import CFA, RawToRgb, RawToRgb2x2Downscaled, RgbToRaw, raw_to_rgb, raw_to_rgb_2x2_downscaled, rgb_to_raw
 from .rgb import (
     BgrToRgb,
     BgrToRgba,
@@ -90,8 +90,10 @@
     "RgbToLab",
     "RgbToRaw",
     "RawToRgb",
+    "RawToRgb2x2Downscaled",
     "raw_to_rgb",
     "rgb_to_raw",
+    "raw_to_rgb_2x2_downscaled",
     "CFA",
     "GrayscaleToRgb",
     "luv_to_rgb",
diff --git a/kornia/color/raw.py b/kornia/color/raw.py
--- a/kornia/color/raw.py
+++ b/kornia/color/raw.py
@@ -3,6 +3,9 @@
 import torch
 from torch import nn
 
+from kornia.core import Module, Tensor, concatenate
+from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_SHAPE
+
 
 class CFA(Enum):
     r"""Define the configuration of the color filter array.
@@ -211,6 +214,65 @@ def rgb_to_raw(image: torch.Tensor, cfa: CFA) -> torch.Tensor:
     return output
 
 
+def raw_to_rgb_2x2_downscaled(image: Tensor, cfa: CFA) -> Tensor:
+    r"""Convert the raw bayer image to RGB version of it and resize width and height by half.
+
+    This is done efficiently by converting each superpixel of bayer image to the corresponding rgb triplet.
+    R and B channels of the raw image are left as are, while two G channels of raw image are averaged to obtain the
+    output G channel.
+
+    We are assuming a CFA with 2 green, 1 red, 1 blue.
+    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 and downscaled with shape :math:`(*,1,H,W)`.
+        cfa: The configuration of the color filter.
+
+    Returns:
+        downscaled RGB version of the image with shape :math:`(*,3,\frac{H}{2},\frac{W}{2})`.
+
+    Example:
+        >>> rawinput = torch.randn(2, 1, 4, 6)
+        >>> rgb = raw_to_rgb_2x2_downscaled(rawinput, CFA.RG) # 2x3x2x3
+    """
+    KORNIA_CHECK(isinstance(image, Tensor), "Input type is not a torch.Tensor")
+
+    KORNIA_CHECK_SHAPE(image, ["*", "1", "H", "W"])
+
+    KORNIA_CHECK(
+        image.shape[-2] % 2 == 0 and image.shape[-1] % 2 == 0,
+        f"Input H&W must be evenly disible by 2. Got {image.shape}",
+    )
+
+    if cfa == CFA.BG:
+        r = image[..., :, ::2, ::2]
+        b = image[..., :, 1::2, 1::2]
+        g1 = image[..., :, ::2, 1::2]
+        g2 = image[..., :, 1::2, ::2]
+    elif cfa == CFA.GB:
+        r = image[..., :, ::2, 1::2]
+        b = image[..., :, 1::2, ::2]
+        g1 = image[..., :, ::2, ::2]
+        g2 = image[..., :, 1::2, 1::2]
+    elif cfa == CFA.RG:
+        r = image[..., :, 1::2, 1::2]
+        b = image[..., :, ::2, ::2]
+        g1 = image[..., :, 1::2, ::2]
+        g2 = image[..., :, ::2, 1::2]
+    elif cfa == CFA.GR:
+        r = image[..., :, 1::2, ::2]
+        b = image[..., :, ::2, 1::2]
+        g1 = image[..., :, 1::2, 1::2]
+        g2 = image[..., :, ::2, ::2]
+    else:
+        raise ValueError(f"Unsupported CFA Got {cfa}.")
+
+    rgb: Tensor = concatenate([r, (g1 + g2) / 2, b], dim=-3)
+
+    return rgb
+
+
 class RawToRgb(nn.Module):
     r"""Module to convert a bayer raw image to RGB version of image.
 
@@ -258,3 +320,27 @@ def __init__(self, cfa: CFA) -> None:
 
     def forward(self, image: torch.Tensor) -> torch.Tensor:
         return rgb_to_raw(image, cfa=self.cfa)
+
+
+class RawToRgb2x2Downscaled(Module):
+    r"""Module version of the :func:`raw_to_rgb_2x2_downscaled()` function.
+
+    The image width and height have to be divisible by two. The image
+    data is assumed to be in the range of (0, 1).
+
+    Shape:
+        - image: :math:`(*, 1, H, W)`
+        - output: :math:`(*, 3, \frac{H}{2}, \frac{W}{2})`
+
+    Example:
+        >>> rawinput = torch.rand(2, 1, 4, 6)
+        >>> rgb_downscale = RawToRgb2x2Downscaled(CFA.RG)
+        >>> output = rgb_downscale(rawinput)  # 2x3x2x3
+    """
+
+    def __init__(self, cfa: CFA) -> None:
+        super().__init__()
+        self.cfa = cfa
+
+    def forward(self, image: Tensor) -> Tensor:
+        return raw_to_rgb_2x2_downscaled(image, cfa=self.cfa)


Test patch:
diff --git a/test/color/test_raw.py b/test/color/test_raw.py
--- a/test/color/test_raw.py
+++ b/test/color/test_raw.py
@@ -196,3 +196,98 @@ def test_module(self, device, 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))
+
+
+class TestRawToRgb2x2Downscaled(BaseTester):
+    def test_smoke(self, device, dtype) -> None:
+        C, H, W = 1, 4, 6
+        img = torch.rand(C, H, W, device=device, dtype=dtype)
+        assert isinstance(kornia.color.raw_to_rgb_2x2_downscaled(img, kornia.color.CFA.BG), torch.Tensor)
+
+    @pytest.mark.parametrize(
+        "batch_size, height, width",
+        [((1,), 6, 4), ((), 6, 4), ((), 0, 0), ((4, 5), 0, 0), ((2,), 2, 4), ((3,), 4, 2), ((4, 2, 1), 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_2x2_downscaled(img, kornia.color.CFA.BG).shape == (
+            *batch_size,
+            3,
+            height // 2,
+            width // 2,
+        )
+
+    def test_exception(self, device, dtype):
+        with pytest.raises(Exception) as errinf:
+            kornia.color.raw_to_rgb_2x2_downscaled([0.0], kornia.color.CFA.BG)
+        assert 'Input type is not a torch.Tensor' in str(errinf)
+
+        with pytest.raises(TypeError) as errinf:
+            img = torch.ones(1, 1, device=device, dtype=dtype)
+            kornia.color.raw_to_rgb_2x2_downscaled(img, kornia.color.CFA.GB)
+        assert "shape must be" in str(errinf)
+
+        with pytest.raises(TypeError) as errinf:
+            img = torch.ones(2, 2, 2, device=device, dtype=dtype)
+            kornia.color.raw_to_rgb_2x2_downscaled(img, kornia.color.CFA.RG)
+        assert "shape must be" in str(errinf)
+
+        with pytest.raises(Exception) as errinf:
+            img = torch.ones(1, 3, 2, device=device, dtype=dtype)
+            kornia.color.raw_to_rgb_2x2_downscaled(img, kornia.color.CFA.GR)
+        assert 'Input H&W must be evenly disible by 2. Got' in str(errinf)
+
+        with pytest.raises(Exception) as errinf:
+            img = torch.ones(1, 2, 3, device=device, dtype=dtype)
+            kornia.color.raw_to_rgb_2x2_downscaled(img, kornia.color.CFA.GR)
+        assert 'Input H&W must be evenly disible by 2. Got' in str(errinf)
+
+        with pytest.raises(ValueError) as errinf:
+            img = torch.ones(1, 4, 8, device=device, dtype=dtype)
+            nonexistent_cfa = 195162495283
+            kornia.color.raw_to_rgb_2x2_downscaled(img, nonexistent_cfa)
+        assert 'Unsupported CFA Got' in str(errinf)
+
+    @pytest.mark.parametrize(
+        "cfa, expected_rgb",
+        [
+            (kornia.color.raw.CFA.BG, [[[0.12, 0.46]], [[0.33, 0.505]], [[0.00, 0.89]]]),
+            (kornia.color.raw.CFA.GB, [[[0.43, 0.58]], [[0.06, 0.675]], [[0.23, 0.43]]]),
+            (kornia.color.raw.CFA.RG, [[[0.0, 0.89]], [[0.33, 0.505]], [[0.12, 0.46]]]),
+            (kornia.color.raw.CFA.GR, [[[0.23, 0.43]], [[0.06, 0.675]], [[0.43, 0.58]]]),
+        ],
+    )
+    def test_functional(self, device, dtype, cfa, expected_rgb):
+        data = torch.tensor([[[0.12, 0.43, 0.46, 0.58], [0.23, 0.00, 0.43, 0.89]]], device=device, dtype=dtype)
+        expected_rgb = torch.tensor(expected_rgb, device=device, dtype=dtype)
+
+        img_rgb = kornia.color.raw_to_rgb_2x2_downscaled(data, cfa)
+        self.assert_close(img_rgb, expected_rgb)
+
+    def test_gradcheck(self, device, dtype):
+        B, C, H, W = 2, 1, 8, 8
+        img = torch.ones(B, C, H, W, device=device, dtype=torch.float64, requires_grad=True)
+        self.gradcheck(
+            kornia.color.raw_to_rgb_2x2_downscaled, (img, kornia.color.raw.CFA.BG), raise_exception=True, fast_mode=True
+        )
+
+    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.color.raw_to_rgb_2x2_downscaled
+        op_jit = torch.jit.script(op)
+        self.assert_close(op(img, kornia.color.raw.CFA.BG), op_jit(img, kornia.color.raw.CFA.BG))
+
+    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.RawToRgb2x2Downscaled(kornia.color.raw.CFA.BG).to(device, dtype)
+        raw_fcn = kornia.color.raw_to_rgb_2x2_downscaled
+        self.assert_close(raw_ops(img), raw_fcn(img, kornia.color.raw.CFA.BG))
+
+    def test_dynamo(self, device, dtype, torch_optimizer):
+        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_2x2_downscaled
+        op_optimized = torch_optimizer(op)
+        self.assert_close(op(img, kornia.color.CFA.BG), op_optimized(img, kornia.color.CFA.BG))