Repository: opencv/opencv
File: modules/dnn/src/layers/nonzero_layer.cpp (C++)

Code before:
        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 depth = CV_MAT_DEPTH(X.type());
        std::vector<size_t> nzcounts;
        std::vector<size_t> nzstart;
        size_t nnz = 0;
        switch (depth)
        {
        case CV_8U:  { nnz = computeNonZeroCountsPerStripe<uchar>(X.ptr<uchar>(), total, nzcounts); break; }
        case CV_8S:  { nnz = computeNonZeroCountsPerStripe<schar>(X.ptr<schar>(), total, nzcounts); break; }
        case CV_16U: { nnz = computeNonZeroCountsPerStripe<ushort>(X.ptr<ushort>(), total, nzcounts); break; }
        case CV_16S: { nnz = computeNonZeroCountsPerStripe<short>(X.ptr<short>(), total, nzcounts); break; }
        case CV_32U: { nnz = computeNonZeroCountsPerStripe<uint32_t>(X.ptr<uint32_t>(), total, nzcounts); break; }
        case CV_32S: { nnz = computeNonZeroCountsPerStripe<int>(X.ptr<int>(), total, nzcounts); break; }
        case CV_64U: { nnz = computeNonZeroCountsPerStripe<uint64_t>(X.ptr<uint64_t>(), total, nzcounts); break; }
        case CV_64S: { nnz = computeNonZeroCountsPerStripe<int64>(X.ptr<int64>(), total, nzcounts); break; }
        case CV_32F: { nnz = computeNonZeroCountsPerStripe<float>(X.ptr<float>(), total, nzcounts); break; }
        case CV_64F: { nnz = computeNonZeroCountsPerStripe<double>(X.ptr<double>(), total, nzcounts); break; }
        case CV_16F: { nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nzcounts, [](uint16_t v){ return isNonZeroF16(v); }); break; }
        case CV_16BF:{ nnz = computeNonZeroCountsPerStripe<uint16_t>(X.ptr<uint16_t>(), total, nzcounts, [](uint16_t v){ return isNonZeroF16(v); }); break; }
        case CV_Bool:{ nnz = computeNonZeroCountsPerStripe<uchar>(X.ptr<uchar>(), total, nzcounts); 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();
        if (kind == _InputArray::STD_VECTOR_MAT) {
            std::vector<Mat>& outs_m = out_arr.getMatVecRef();
            outs_m.resize(1);
            outs_m[0].fit(outShape, CV_64S);
            Mat& Y = outs_m[0];
            if (nnz == 0) return;

Review comment (bug):
```
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);
}
int64* y = Y.ptr<int64>();

switch (depth) {
      ....
}

if (kind == _InputArray::STD_VECTOR_UMAT) {
     Y.copyTo(out_umats->at(0));
}
```

Code after:
        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; }