From 57909cccc6eb7ff544c91734462fd7b1e44f8322 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Zientkiewicz?= <michalz@nvidia.com>
Date: Fri, 27 Sep 2024 15:29:49 +0200
Subject: [PATCH] Fix CPU resize with mixed NN/other resampling filters.
 (#5647)

* Fix size computation (handle negative support size)
* Add regression tests

---------

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
---
 dali/kernels/imgproc/resample/separable_cpu.h |  4 +-
 dali/test/python/operator_2/test_resize.py    | 46 +++++++++++++++++++
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/dali/kernels/imgproc/resample/separable_cpu.h b/dali/kernels/imgproc/resample/separable_cpu.h
index 0ccfd713b87..1771e67ffcf 100644
--- a/dali/kernels/imgproc/resample/separable_cpu.h
+++ b/dali/kernels/imgproc/resample/separable_cpu.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2019-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2019-2021, 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -96,7 +96,7 @@ struct ResamplingSetupCPU : SeparableResamplingSetup<_spatial_ndim> {
 
       for (int stage = 0; stage < num_stages; stage++) {
         size_t extent = resized_dim_extent(stage);
-        size_t support = filter_support(stage);
+        int support = std::max(0, filter_support(stage));  // support() returns -1 for NN interp.
         size_t num_coeffs = extent * support;
 
         if (extent > req.indices_size)
diff --git a/dali/test/python/operator_2/test_resize.py b/dali/test/python/operator_2/test_resize.py
index 11280599dac..fc4e05a6442 100644
--- a/dali/test/python/operator_2/test_resize.py
+++ b/dali/test/python/operator_2/test_resize.py
@@ -843,6 +843,52 @@ def resize_pipe():
     assert np.max(np.abs(out - large_data_resized)) < 2
 
 
+@params(("cpu", 0), ("cpu", 1), ("gpu", 0), ("gpu", 1))
+def test_nn_on_one_axis(device, axis):
+    # Checks whether having NN interpolation in one axis and full resampling in the other works
+    data = np.array(
+        [
+            [0, 0, 0],
+            [0, 1, 0],
+            [0, 2, 0],
+        ],
+        dtype=np.float32,
+    )
+
+    # magnification is NN, minification is triangular
+    ref = np.array(
+        [
+            [0, 0, 0.3333334, 0.3333334, 0, 0],
+            [0, 0, 1.6666667, 1.6666667, 0, 0],
+        ],
+        dtype=np.float32,
+    )
+
+    if axis == 1:
+        data = np.transpose(data, (1, 0))
+        ref = np.transpose(ref, (1, 0))
+
+    # add channel
+    data = data[..., np.newaxis]
+    ref = ref[..., np.newaxis]
+
+    @pipeline_def(batch_size=1, device_id=0, num_threads=1)
+    def test_pipe():
+        src = dali.types.Constant(data, device=device)
+        return fn.resize(
+            src,
+            size=ref.shape[:-1],
+            min_filter=dali.types.INTERP_LINEAR,
+            mag_filter=dali.types.INTERP_NN,
+            antialias=True,
+        )
+
+    pipe = test_pipe()
+    pipe.build()
+    (out,) = pipe.run()
+    check_batch(out, [ref], 1, 1e-5, 1e-5, None, False)
+
+
 def test_checkerboard_dali_vs_onnx_ref():
     improc_data_dir = os.path.join(test_data_root, "db", "imgproc")
     ref_dir = os.path.join(improc_data_dir, "ref", "resampling")