Merge pull request #27701 from abhishek-gola:nonzero_layer_add

Added nonzero layer to new DNN engine #27701

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake

modules/dnn/include/opencv2/dnn/all_layers.hpp

@@ -1374,6 +1374,12 @@ CV__DNN_INLINE_NS_BEGIN
         static Ptr<Expand2Layer> create(const LayerParams &params);
     };
 
+    class CV_EXPORTS NonZeroLayer : public Layer
+    {
+    public:
+        static Ptr<NonZeroLayer> create(const LayerParams& params);
+    };
+
     class CV_EXPORTS InstanceNormLayer : public Layer {
     public:
         float epsilon;

modules/dnn/src/init.cpp

@@ -94,6 +94,7 @@ void initializeLayerFactory()
     CV_DNN_REGISTER_LAYER_CLASS(Flatten,        FlattenLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Interp,         InterpLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Pad2,           Pad2Layer);
+    CV_DNN_REGISTER_LAYER_CLASS(NonZero,        NonZeroLayer);
     CV_DNN_REGISTER_LAYER_CLASS(QuantizeLinear, QuantizeLinearLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Range,          RangeLayer);
     CV_DNN_REGISTER_LAYER_CLASS(Reshape,        ReshapeLayer);

modules/dnn/src/layers/nonzero_layer.cpp

@@ -0,0 +1,243 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the
+// top-level directory of this distribution and at http://opencv.org/license.html.
+
+#include "../precomp.hpp"
+#include "layers_common.hpp"
+#include <opencv2/dnn/shape_utils.hpp>
+#include <cstdint>
+#include <array>
+
+namespace cv {
+namespace dnn {
+
+// ONNX NonZero operator
+// Spec: https://onnx.ai/onnx/operators/onnx__NonZero.html
+// Supported opsets: 9-18 (no attributes)
+
+namespace {
+static inline bool isNonZeroF16(uint16_t val)
+{
+    return (val & 0x7fff) != 0;
+}
+
+template <typename T, typename PT>
+static size_t computeNonZeroCountsPerStripe(const T* data, size_t total, int nstripes, std::vector<size_t>& nzcounts, PT isNonZero)
+{
+    nzcounts.assign(nstripes, 0);
+    if (total == 0)
+        return 0;
+    cv::parallel_for_(cv::Range(0, nstripes), [&](const cv::Range& range) {
+        const size_t i0 = total * (size_t)range.start / (size_t)nstripes;
+        const size_t i1 = total * (size_t)range.end   / (size_t)nstripes;
+        size_t nz = 0;
+        for (size_t i = i0; i < i1; ++i)
+            nz += isNonZero(data[i]);
+        nzcounts[(size_t)range.start] = nz;
+    });
+
+    size_t total_nz = 0;
+    for (size_t c : nzcounts) total_nz += c;
+    return total_nz;
+}
+
+template <typename T>
+static inline size_t computeNonZeroCountsPerStripe(const T* data, size_t total, int nstripes, std::vector<size_t>& nzcounts)
+{
+    return computeNonZeroCountsPerStripe<T>(data, total, nstripes, nzcounts, [](T v){ return v != (T)0; });
+}
+
+template <typename T, typename PT>
+static void emitIndicesStripes(const T* data,
+                               size_t total,
+                               int nstripes,
+                               int rank,
+                               const std::vector<int>& dims,
+                               const std::vector<int>& strides,
+                               const std::vector<size_t>& nzstart,
+                               int64* y,
+                               PT isNonZero)
+{
+    const size_t nnz_total = nzstart.back();
+    if (rank == 0) {
+        return;
+    }
+    const int lastDimSize = dims[rank - 1];
+
+    cv::parallel_for_(cv::Range(0, nstripes), [&](const cv::Range& range) {
+        size_t i = total * (size_t)range.start / (size_t)nstripes;
+        const size_t i1 = total * (size_t)range.end / (size_t)nstripes;
+        size_t col = nzstart[(size_t)range.start];
+
+        std::array<int, CV_MAX_DIM> coord;
+        std::fill_n(coord.begin(), rank, 0);
+        if (rank > 0) coord[rank - 1] = lastDimSize - 1;
+
+        for (; i < i1; ++i)
+        {
+            if (rank > 0 && ++coord[rank - 1] >= lastDimSize)
+            {
+                size_t idxLinear = i;
+                for (int d = rank - 1; d >= 0; --d)
+                {
+                    coord[d] = static_cast<int>(idxLinear % (size_t)dims[d]);
+                    idxLinear /= (size_t)dims[d];
+                }
+            }
+            if (isNonZero(data[i]))
+            {
+                for (int d = 0; d < rank; ++d)
+                    y[(size_t)d * nnz_total + col] = (int64)coord[d];
+                ++col;
+            }
+        }
+    });
+}
+
+template <typename T>
+static void emitIndicesStripes(const T* data,
+                               size_t total,
+                               int nstripes,
+                               int rank,
+                               const std::vector<int>& dims,
+                               const std::vector<int>& strides,
+                               const std::vector<size_t>& nzstart,
+                               int64* y)
+{
+    emitIndicesStripes<T>(data, total, nstripes, rank, dims, strides, nzstart, y, [](T v){ return v != (T)0; });
+}
+}
+
+class NonZeroLayerImpl CV_FINAL : public NonZeroLayer
+{
+public:
+    NonZeroLayerImpl(const LayerParams& params) { setParamsFrom(params); }
+
+    virtual bool dynamicOutputShapes() const CV_OVERRIDE
+    {
+        return true;
+    }
+
+    bool supportBackend(int backendId) CV_OVERRIDE
+    {
+        return backendId == DNN_BACKEND_OPENCV;
+    }
+
+    bool getMemoryShapes(const std::vector<MatShape>& inputs, const int /*requiredOutputs*/,
+                         std::vector<MatShape>& outputs, std::vector<MatShape>& /*internals*/) const CV_OVERRIDE
+    {
+        CV_Assert(inputs.size() == 1);
+        const MatShape& in = inputs[0];
+        CV_Assert(!in.empty());
+
+        int rank = (int)in.size();
+        MatShape out = shape(rank, -1);
+        outputs.assign(1, out);
+        return false;
+    }
+
+    void getTypes(const std::vector<MatType>& /*inputs*/, const int requiredOutputs,
+                  const int requiredInternals, std::vector<MatType>& outputs,
+                  std::vector<MatType>& internals) const CV_OVERRIDE
+    {
+        outputs.assign(requiredOutputs, CV_64S);
+    }
+
+    void forward(InputArrayOfArrays in_arr, OutputArrayOfArrays out_arr, OutputArrayOfArrays) CV_OVERRIDE
+    {
+        CV_Assert(in_arr.size().area() == 1);
+        const Mat X = in_arr.getMat(0);
+        CV_Assert(X.data != nullptr);
+
+        const int rank = X.dims;
+        std::vector<int> dims(rank), strides(rank);
+        for (int i = 0; i < rank; ++i) dims[i] = X.size[i];
+
+        if (rank > 0) {
+            strides.back() = 1;
+            for (int i = rank - 2; i >= 0; --i)
+                strides[i] = strides[i + 1] * dims[i + 1];
+        }
+
+        size_t total = X.total();
+        const int nstripes = 16;
+
+        const int depth = CV_MAT_DEPTH(X.type());
+        std::vector<size_t> nzcounts;
+        std::vector<size_t> nzstart;
+        size_t nnz = 0;
+        switch (depth)
+        {
+        case CV_Bool:
+        case CV_8U:
+        case CV_8S:  { nnz = computeNonZeroCountsPerStripe<uchar>(X.ptr<uchar>(), total, nstripes, nzcounts); break; }
+        case CV_16U:
+        case CV_16S: { nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nstripes, nzcounts); break; }
+        case CV_32U:
+        case CV_32S: { nnz = computeNonZeroCountsPerStripe<uint32_t>(X.ptr<uint32_t>(), total, nstripes, nzcounts); break; }
+        case CV_64U:
+        case CV_64S: { nnz = computeNonZeroCountsPerStripe<uint64_t>(X.ptr<uint64_t>(), total, nstripes, nzcounts); break; }
+        case CV_32F: { nnz = computeNonZeroCountsPerStripe<float>(X.ptr<float>(), total, nstripes, nzcounts); break; }
+        case CV_64F: { nnz = computeNonZeroCountsPerStripe<double>(X.ptr<double>(), total, nstripes, nzcounts); break; }
+        case CV_16F:
+        case CV_16BF:{ nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nstripes, nzcounts, [](uint16_t v){ return isNonZeroF16(v); }); break; }
+        default:
+            CV_Error_(Error::StsError, ("NonZero: Unsupported input depth=%d", depth));
+        }
+
+        nzstart.assign(nzcounts.size() + 1, 0);
+        for (size_t i = 1; i <= nzcounts.size(); ++i)
+            nzstart[i] = nzstart[i - 1] + nzcounts[i - 1];
+
+        MatShape outShape = shape(rank, (int)nnz);
+        auto kind = out_arr.kind();
+        std::vector<Mat>* out_mats = nullptr;
+        std::vector<UMat>* out_umats = nullptr;
+        Mat Y;
+        if (kind == _InputArray::STD_VECTOR_MAT) {
+            out_mats = &out_arr.getMatVecRef();
+            out_mats->resize(1);
+            out_mats->at(0).fit(outShape, CV_64S);
+            Y = out_mats->at(0);
+        } else {
+            CV_Assert(kind == _InputArray::STD_VECTOR_UMAT);
+            out_umats = &out_arr.getUMatVecRef();
+            out_umats->resize(1);
+            out_umats->at(0).fit(outShape, CV_64S);
+            Y = Mat(outShape, CV_64S);
+        }
+
+        if (nnz == 0) return;
+
+        int64* y = Y.ptr<int64>();
+
+        switch (depth)
+        {
+        case CV_Bool:
+        case CV_8U:
+        case CV_8S:  { emitIndicesStripes<uchar>(X.ptr<uchar>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_16U:
+        case CV_16S: { emitIndicesStripes<uint16_t>(X.ptr<uint16_t>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_32U:
+        case CV_32S: { emitIndicesStripes<uint32_t>(X.ptr<uint32_t>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_64U:
+        case CV_64S: { emitIndicesStripes<uint64_t>(X.ptr<uint64_t>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_32F: { emitIndicesStripes<float>(X.ptr<float>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_64F: { emitIndicesStripes<double>(X.ptr<double>(), total, nstripes, rank, dims, strides, nzstart, y); break; }
+        case CV_16F:
+        case CV_16BF:{ emitIndicesStripes<uint16_t>(X.ptr<uint16_t>(), total, nstripes, rank, dims, strides, nzstart, y, [](uint16_t v){ return isNonZeroF16(v); }); break; }
+        default: break;
+        }
+
+        if (kind == _InputArray::STD_VECTOR_UMAT) {
+            Y.copyTo(out_umats->at(0));
+        }
+    }
+};
+
+Ptr<NonZeroLayer> NonZeroLayer::create(const LayerParams& params)
+{
+    return makePtr<NonZeroLayerImpl>(params);
+}
+
+}} // namespace cv::dnn

modules/dnn/src/onnx/onnx_importer2.cpp

@@ -208,6 +208,7 @@ protected:
     void parsePRelu                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseRange                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseReduce               (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
+    void parseNonZero              (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseRelu                 (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseTrilu                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
     void parseIsNaN                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
@@ -1712,6 +1713,12 @@ void ONNXImporter2::parseUpsample(LayerParams& layerParams, const opencv_onnx::N
     addLayer(layerParams, node_proto, n_inputs);
 }
 
+void ONNXImporter2::parseNonZero(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
+{
+    layerParams.type = "NonZero";
+    addLayer(layerParams, node_proto);
+}
+
 void ONNXImporter2::parseSoftMax(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
 {
     const std::string& layer_type = node_proto.op_type();
@@ -2465,6 +2472,7 @@ void ONNXImporter2::buildDispatchMap_ONNX_AI(int opset_version)
     dispatch["Tanh"] = &ONNXImporter2::parseTanh;
     dispatch["Abs"] = &ONNXImporter2::parseAbs;
     dispatch["PRelu"] = &ONNXImporter2::parsePRelu;
+    dispatch["NonZero"] = &ONNXImporter2::parseNonZero;
     dispatch["LRN"] = &ONNXImporter2::parseLRN;
     dispatch["InstanceNormalization"] = &ONNXImporter2::parseInstanceNormalization;
     dispatch["BatchNormalization"] = &ONNXImporter2::parseBatchNormalization;

modules/dnn/test/test_onnx_conformance_layer_filter__openvino.inl.hpp

@@ -1218,7 +1218,7 @@ CASE(test_nonmaxsuppression_two_batches)
 CASE(test_nonmaxsuppression_two_classes)
     // no filter
 CASE(test_nonzero_example)
-    // no filter
+    SKIP;
 CASE(test_not_2d)
     // no filter
 CASE(test_not_3d)

modules/dnn/test/test_onnx_conformance_layer_filter_opencv_classic_denylist.inl.hpp

@@ -139,3 +139,4 @@
 "test_constant_pad",
 "test_constantofshape_float_ones",
 "test_constantofshape_int_zeros",
+"test_nonzero_example",

modules/dnn/test/test_onnx_conformance_layer_parser_denylist.inl.hpp

@@ -132,7 +132,6 @@
 "test_nonmaxsuppression_suppress_by_IOU_and_scores", // ---- same as above ---
 "test_nonmaxsuppression_two_batches", // ---- same as above ---
 "test_nonmaxsuppression_two_classes", // ---- same as above ---
-"test_nonzero_example", // Issue::Can't create layer "onnx_node_output_0!result" of type "NonZero" in function 'getLayerInstance'
 "test_onehot_negative_indices", // Issue:: Layer does not exist (OneHot) :: Can't create layer "onnx_node_output_0!y" of type "OneHot" in function 'getLayerInstance'
 "test_onehot_with_axis", // ---- same as above ---
 "test_onehot_with_negative_axis", // ---- same as above ---