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

Code before:
        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;

            int64* y = Y.ptr<int64>();

            switch (depth)
            {
            case CV_8U:  { emitIndicesStripes<uchar>(X.ptr<uchar>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_8S:  { emitIndicesStripes<schar>(X.ptr<schar>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_16U: { emitIndicesStripes<ushort>(X.ptr<ushort>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_16S: { emitIndicesStripes<short>(X.ptr<short>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_32U: { emitIndicesStripes<uint32_t>(X.ptr<uint32_t>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_32S: { emitIndicesStripes<int>(X.ptr<int>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_64U: { emitIndicesStripes<uint64_t>(X.ptr<uint64_t>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_64S: { emitIndicesStripes<int64>(X.ptr<int64>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_32F: { emitIndicesStripes<float>(X.ptr<float>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_64F: { emitIndicesStripes<double>(X.ptr<double>(), total, rank, dims, strides, nzstart, y); break; }
            case CV_16F: { emitIndicesStripes<uint16_t>(X.ptr<uint16_t>(), total, rank, dims, strides, nzstart, y, [](uint16_t v){ return isNonZeroF16(v); }); break; }
            case CV_16BF:{ emitIndicesStripes<uint16_t>(X.ptr<uint16_t>(), total, rank, dims, strides, nzstart, y, [](uint16_t v){ return isNonZeroF16(v); }); break; }
            case CV_Bool:{ emitIndicesStripes<uchar>(X.ptr<uchar>(), total, rank, dims, strides, nzstart, y); break; }
            default: break;
            }
        } else if (kind == _InputArray::STD_VECTOR_UMAT) {
            std::vector<UMat>& outs_u = out_arr.getUMatVecRef();
            outs_u.resize(1);
            outs_u[0].fit(outShape, CV_64S);
            if (nnz == 0) return;

            Mat Y(outShape, CV_64S);
            int64* y = Y.ptr<int64>();


Review comment (refactor):
can we retain only a single switch statement? All the memory allocation logic can be put into helper methods, that we will call from each branch of this switch.

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