Add tests for opencv cpu algos. Bugfix output channel for grey decode…

…. Add assert and note for vng only supports uint8.

Signed-off-by: Bryce Ferenczi <[email protected]>

dali/operators/image/color/debayer.cc

@@ -101,13 +101,13 @@ Different algorithms are supported on the GPU and CPU.
 
 **GPU Algorithms:**
 
- - ``bilinear_npp`` - - uses bilinear interpolation with chroma correlation for green values.
+ - ``bilinear_npp`` - - bilinear interpolation with chroma correlation for green values.
 
 **CPU Algorithms:**
 
- - ``bilinear_ocv`` - uses bilinear interpolation.
- - ``edgeaware_ocv`` - uses edge-aware interpolation.
- - ``vng_ocv`` - uses Variable Number of Gradients (VNG) interpolation.
+ - ``bilinear_ocv`` - bilinear interpolation.
+ - ``edgeaware_ocv`` - edge-aware interpolation.
+ - ``vng_ocv`` - Variable Number of Gradients (VNG) interpolation (only ``uint8_t`` supported).
  - ``gray_ocv`` - converts the image to grayscale with bilinear interpolation.)code",
             DALI_STRING)
     .InputLayout(0, {"HW", "HWC", "FHW", "FHWC"})

dali/operators/image/color/debayer_cpu.cc

@@ -95,6 +95,10 @@ class DebayerCPU : public Debayer<CPUBackend> {
   bool SetupImpl(std::vector<OutputDesc> &output_desc, const Workspace &ws) {
     DALI_ENFORCE(alg_ != debayer::DALIDebayerAlgorithm::DALI_DEBAYER_BILINEAR_NPP,
                  "bilinear_npp algorithm is not supported on CPU.");
+    if (alg_ == debayer::DALIDebayerAlgorithm::DALI_DEBAYER_VNG_OCV) {
+      DALI_ENFORCE(ws.Input<CPUBackend>(0).type() == DALI_UINT8,
+                   "VNG debayering only supported with UINT8.");
+    }
     return Debayer<CPUBackend>::SetupImpl(output_desc, ws);
   }
 
@@ -109,12 +113,13 @@ class DebayerCPU : public Debayer<CPUBackend> {
         const auto inImage = input[i];
         auto outImage = output[i];
 
-        const auto &inShape = inImage.shape();
-        const auto height = static_cast<int>(inShape[0]);
-        const auto width = static_cast<int>(inShape[1]);
+        const auto &oShape = outImage.shape();
+        const auto height = static_cast<int>(oShape[0]);
+        const auto width = static_cast<int>(oShape[1]);
+        const auto outChannels = static_cast<int>(oShape[2]);
         cv::Mat inImg(height, width, OCVMatTypeForDALIData(inImage.type(), 1),
                       const_cast<void *>(inImage.raw_data()));
-        cv::Mat outImg(height, width, OCVMatTypeForDALIData(outImage.type(), 3),
+        cv::Mat outImg(height, width, OCVMatTypeForDALIData(outImage.type(), outChannels),
                        outImage.raw_mutable_data());
         cv::demosaicing(inImg, outImg, toOpenCVColorConversionCode(pattern_[i], alg_));
       });

dali/test/python/debayer_test_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import cv2
 import numpy as np
 from scipy.signal import convolve2d
 
@@ -177,3 +178,18 @@ def debayer_bilinear_npp_pattern_seq(seq, patterns):
     seq_masks = [bayer_masks[pattern] for pattern in patterns]
     reds, greens, blues = [np.stack(channel) for channel in zip(*seq_masks)]
     return debayer_bilinear_npp_masks(seq, (reds, greens, blues))
+
+
+def debayer_opencv(img: np.ndarray, pattern: BayerPattern, algorithm: str):
+    """Debayer image with OpenCV."""
+    if pattern is not BayerPattern.BGGR:
+        raise NotImplementedError("Only BGGR pattern is supported by at the moment.")
+    if algorithm == "bilinear_ocv":
+        return cv2.cvtColor(img, cv2.COLOR_BayerBG2RGB)
+    if algorithm == "edgeaware_ocv":
+        return cv2.cvtColor(img, cv2.COLOR_BayerBG2RGB_EA)
+    if algorithm == "vng_ocv":
+        return cv2.cvtColor(img, cv2.COLOR_BayerBG2RGB_VNG)
+    if algorithm == "gray_ocv":
+        return cv2.cvtColor(img, cv2.COLOR_BayerBG2GRAY)[..., None]  # Make HWC
+    raise ValueError(f"Unknown algorithm: {algorithm}")

dali/test/python/operator_1/test_debayer.py

@@ -17,20 +17,21 @@
 import unittest
 
 import numpy as np
-
-from nvidia.dali import pipeline_def, fn, types
-from test_utils import get_dali_extra_path
-from nose_utils import assert_raises
-from nose2.tools import params, cartesian_params
 from debayer_test_utils import (
+    BayerPattern,
     bayer_patterns,
     blue_position,
     blue_position2pattern,
-    rgb2bayer,
-    rgb2bayer_seq,
     debayer_bilinear_npp_pattern,
     debayer_bilinear_npp_pattern_seq,
+    debayer_opencv,
+    rgb2bayer,
+    rgb2bayer_seq,
 )
+from nose2.tools import cartesian_params, params
+from nose_utils import assert_raises
+from nvidia.dali import fn, pipeline_def, types
+from test_utils import get_dali_extra_path
 
 data_root = get_dali_extra_path()
 images_dir = os.path.join(data_root, "db", "single", "jpeg")
@@ -227,6 +228,52 @@ def debayer_pipeline():
                 baseline = npp_baseline[pattern][idx]
                 assert compare_image_equality(img_debayered, baseline, device)
 
+    @cartesian_params(
+        ("bilinear_ocv", "edgeaware_ocv", "vng_ocv", "gray_ocv"), (np.uint8, np.uint16)
+    )
+    def test_cpu_algorithms(self, algorithm: str, dtype: np.dtype):
+        if algorithm == "vng_ocv" and dtype == np.uint16:
+            # VNG algorithm is not supported for uint16
+            return
+
+        num_iterations = 3
+        batch_size = 1
+        bayered_imgs, _ = self.get_test_data(dtype)
+
+        pattern = BayerPattern.BGGR
+
+        def source(sample_info):
+            idx = sample_info.idx_in_epoch % self.num_samples
+            return (
+                bayered_imgs[pattern][idx],
+                np.array(idx, dtype=np.int32),
+            )
+
+        @pipeline_def
+        def debayer_pipeline():
+            bayer_imgs, idxs = fn.external_source(source=source, batch=False, num_outputs=2)
+            debayered_imgs = fn.experimental.debayer(
+                bayer_imgs, blue_position=blue_position(pattern), algorithm=algorithm
+            )
+            return debayered_imgs, idxs
+
+        pipe = debayer_pipeline(batch_size=batch_size, device_id=0, num_threads=4)
+
+        out_batches = []
+        for _ in range(num_iterations):
+            debayered_imgs_dev, idxs = pipe.run()
+            assert debayered_imgs_dev.layout() == "HWC"
+            out_batches.append(
+                (list(map(np.array, debayered_imgs_dev.as_cpu())), list(map(np.array, idxs)))
+            )
+
+        for debayered_imgs, idxs in out_batches:
+            assert len(debayered_imgs) == len(idxs)
+            for img_debayered, idx in zip(debayered_imgs, idxs):
+                baseline = bayered_imgs[pattern][idx]
+                baseline = debayer_opencv(baseline, pattern, algorithm)
+                assert np.all(img_debayered == baseline)
+
 
 class DebayerVideoTest(unittest.TestCase):
     @classmethod