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

Code before:
                for (int d = (int)rankIdx.size() - 1; d >= 0; --d) {
                    int dim = rankIdx[d];
                    int coord = rem % dim;
                    rem /= dim;
                    base += (size_t)coord * mappedStrides[d];
                }
                int64_t idxVal = idxPtr64[i];
                bool wasOutOfRange = (idxVal < 0 || idxVal >= depth);
                if (idxVal < 0) { localNeg++; }
                if (wasOutOfRange) { localOOR++; }
                int64_t wrapped = idxVal % depth;
                if (wrapped < 0) wrapped += depth;
                size_t pos = base + strides[insAxis] * (size_t)wrapped;
                outPtr[pos] = onVal;
                localSet++;
            }
            numSet += localSet;
            numNegWrapped += localNeg;
            numOOR += localOOR;
        });
    }

public:
    void forward(InputArrayOfArrays in, OutputArrayOfArrays out, OutputArrayOfArrays) CV_OVERRIDE
    {
        std::vector<Mat> inputs, outputs; in.getMatVector(inputs); out.getMatVector(outputs);
        CV_Assert(inputs.size() >= 2);
        CV_Assert(outputs.size() == 1);

        const Mat& indices = inputs[0];
        const Mat& depthMat = inputs[1];
        Mat values; if (inputs.size() >= 3) values = inputs[2];
        Mat& Y = outputs[0];

        int outDepth = values.empty() ? CV_32F : values.depth();

        MatShape outShape = shape(indices);
        int ndims = (int)outShape.size();
        int insAxis = normalize_axis(axis, ndims + 1);

        MatShape finalShape;
        if (depthMat.total() == 1)
        {
            int64_t depth64 = 0;
            tensorToScalar(depthMat, CV_64S, &depth64);
            if (depth64 > 0)
            {
                outShape.insert(outShape.begin() + insAxis, (int)depth64);
                finalShape = outShape;
            }
        }

Review comment (suggestion):
since the 'depth' tensor is not necessarily constant, don't use getMatVector. Instead, check the 'kind' and use getMatVecRef/getUMatVecRef, depending on the kind

Code after:
                if (wrapped < 0) wrapped += localDepth;
                size_t hi = (size_t)pos / localOutStep;
                size_t lo = (size_t)pos % localOutStep;
                size_t base = hi * ((size_t)localDepth * localOutStep) + lo;
                localOutPtr[base + (size_t)wrapped * localOutStep] = localOnVal;
            }
        });
    }

public:
    void forward(InputArrayOfArrays in, OutputArrayOfArrays out, OutputArrayOfArrays) CV_OVERRIDE
    {
        int inKind = in.kind();
        int outKind = out.kind();
        CV_Assert(in.size().area() == 3);

        Mat indices, depthMat, values;
        if (inKind == _InputArray::STD_VECTOR_MAT)
        {
            indices = in.getMat(0);
            depthMat = in.getMat(1);
            values = in.getMat(2);
        }
        else if (inKind == _InputArray::STD_VECTOR_UMAT)
        {
            indices = in.getUMat(0).getMat(ACCESS_READ);
            depthMat = in.getUMat(1).getMat(ACCESS_READ);
            values = in.getUMat(2).getMat(ACCESS_READ);
        }
        else
        {
            CV_Error(Error::StsNotImplemented, "OneHot: unsupported input type");
        }

        std::vector<Mat>* out_mats = nullptr;
        std::vector<UMat>* out_umats = nullptr;
        Mat Y;
        if (outKind == _InputArray::STD_VECTOR_MAT)
        {
            out_mats = &out.getMatVecRef();
            out_mats->resize(1);
        }
        else if (outKind == _InputArray::STD_VECTOR_UMAT)
        {
            out_umats = &out.getUMatVecRef();
            out_umats->resize(1);
        }
        else
        {
            CV_Error(Error::StsNotImplemented, "OneHot: unsupported output type");
        }