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

Code before:
        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;
            }
        }
        if (!finalShape.empty())
        {
            outputs[0].fit(finalShape, outDepth);
        }

        switch (outDepth) {
            case CV_8U:   runKernel<uchar>(indices, depthMat, values, Y); break;
            case CV_8S:   runKernel<schar>(indices, depthMat, values, Y); break;
            case CV_16U:  runKernel<uint16_t>(indices, depthMat, values, Y); break;
            case CV_16S:  runKernel<int16_t>(indices, depthMat, values, Y); break;
            case CV_32S:  runKernel<int>(indices, depthMat, values, Y); break;
            case CV_64S:  runKernel<int64_t>(indices, depthMat, values, Y); break;
            case CV_16F:  runKernel<hfloat>(indices, depthMat, values, Y); break;
            case CV_16BF: runKernel<bfloat>(indices, depthMat, values, Y); break;
            case CV_32F:  runKernel<float>(indices, depthMat, values, Y); break;
            case CV_64F:  runKernel<double>(indices, depthMat, values, Y); break;
            default: CV_Error(Error::BadDepth, "OneHot: unsupported output type");
        }
    }
};

Ptr<OneHotLayer> OneHotLayer::create(const LayerParams& params)
{
    return makePtr<OneHotLayerImpl>(params);
}

}} // namespace cv::dnn


Review comment (suggestion):
no need to have different branches for int8/uint8/bool, int16/uint16/float16/bfloat16, int32/uint32/float, int64/uint64/double. Just use a single kernel for all types with the same element size.

Code after:
        }

        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");
        }

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

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

        CV_Assert(depthMat.total() == 1);
        int64_t depth64 = 0;
        tensorToScalar(depthMat, CV_64S, &depth64);
        CV_CheckGT(depth64, 0, "OneHot: depth must be > 0");
        MatShape finalShape;
        outShape.insert(outShape.begin() + insAxis, (int)depth64);
        finalShape = outShape;
        if (!finalShape.empty())
        {
            if (outKind == _InputArray::STD_VECTOR_MAT)
            {
                out_mats->at(0).fit(finalShape, outDepth);
                Y = out_mats->at(0);
            }
            else
            {
                out_umats->at(0).fit(finalShape, outDepth);
                Y = Mat(finalShape, outDepth);
            }
        }

        switch (outDepth) {
            case CV_8U:   runKernel<uchar>(indices, (int)depth64, values, Y); break;
            case CV_8S:   runKernel<schar>(indices, (int)depth64, values, Y); break;
            case CV_16U:  runKernel<uint16_t>(indices, (int)depth64, values, Y); break;