Merge pull request #27902 from abhishek-gola:onehot_layer_add

Added Onehot layer support to new DNN engine #27902 ### 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
2026-07-29 15:23:05 +04:00 · 2025-10-31 00:51:14 +05:30
parent c8eafa3ee4
commit e9298fb73e
7 changed files with 293 additions and 8 deletions
@@ -556,6 +556,13 @@ CV__DNN_INLINE_NS_BEGIN
        static Ptr<GridSampleLayer> create(const LayerParams& params);
    };

+    class CV_EXPORTS OneHotLayer : public Layer
+    {
+    public:
+        int axis;
+        static Ptr<OneHotLayer> create(const LayerParams& params);
+    };
+
    class CV_EXPORTS FlattenLayer : public Layer
    {
    public:
@@ -115,6 +115,7 @@ void initializeLayerFactory()
    CV_DNN_REGISTER_LAYER_CLASS(Unsqueeze,      UnsqueezeLayer);
    CV_DNN_REGISTER_LAYER_CLASS(IsNaN,          IsNaNLayer);
    CV_DNN_REGISTER_LAYER_CLASS(IsInf,          IsInfLayer);
+    CV_DNN_REGISTER_LAYER_CLASS(OneHot,         OneHotLayer);
    CV_DNN_REGISTER_LAYER_CLASS(Det,            DetLayer);
    CV_DNN_REGISTER_LAYER_CLASS(CenterCropPad,  CenterCropPadLayer);
    CV_DNN_REGISTER_LAYER_CLASS(BitShift,       BitShiftLayer);
@@ -0,0 +1,269 @@
+// 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.
+// Copyright (C) 2025, BigVision LLC, all rights reserved.
+// Third party copyrights are property of their respective owners.
+
+#include "../precomp.hpp"
+#include "../net_impl.hpp"
+#include "layers_common.hpp"
+#include <opencv2/dnn/shape_utils.hpp>
+#include <opencv2/core/utils/logger.hpp>
+#include <atomic>
+
+namespace cv { namespace dnn {
+
+/*
+    OneHot layer, as defined in ONNX specification:
+    https://onnx.ai/onnx/operators/onnx__OneHot.html
+    Supported opsets: 9-11
+*/
+
+class OneHotLayerImpl CV_FINAL : public OneHotLayer
+{
+public:
+    OneHotLayerImpl(const LayerParams& params)
+    {
+        setParamsFrom(params);
+        axis = params.get<int>("axis", -1);
+    }
+
+    bool supportBackend(int backendId) CV_OVERRIDE
+    {
+        return backendId == DNN_BACKEND_OPENCV;
+    }
+
+private:
+    bool getConstPositiveDepth(int& depthOut) const
+    {
+        Net::Impl* netimpl_ = getNetImpl(this);
+        if (!netimpl_ || !netimpl_->isConstArg(this->inputs[1]))
+            return false;
+
+        try
+        {
+            Mat depthTensor = netimpl_->argTensor(this->inputs[1]);
+            if (!depthTensor.empty() && depthTensor.total() == 1)
+            {
+                int64_t depth64 = 0;
+                tensorToScalar(depthTensor, CV_64S, &depth64);
+                if (depth64 > 0)
+                {
+                    depthOut = (int)depth64;
+                    return true;
+                }
+            }
+        }
+        catch (const cv::Exception& e)
+        {
+            CV_Error_(Error::StsError, ("OneHot: failed to access constant depth: %s", e.what()));
+        }
+        return false;
+    }
+
+    bool tryInferStaticOutputShape(const MatShape& idxShape, MatShape& outShape) const
+    {
+        outShape.clear();
+        int depth = 0;
+        if (getConstPositiveDepth(depth))
+        {
+            outShape = idxShape;
+            int insAxis = normalize_axis(axis, (int)outShape.size() + 1);
+            outShape.insert(outShape.begin() + insAxis, depth);
+            return true;
+        }
+        return false;
+    }
+
+public:
+    bool dynamicOutputShapes() const CV_OVERRIDE
+    {
+        Net::Impl* netimpl_ = getNetImpl(this);
+        return !(netimpl_ && netimpl_->isConstArg(this->inputs[1]));
+    }
+
+    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() == 3);
+        const MatShape& idxShape = inputs[0];
+        CV_Assert(!idxShape.empty());
+
+        MatShape outShape;
+        if (tryInferStaticOutputShape(idxShape, outShape))
+        {
+            outputs.assign(1, outShape);
+            internals.clear();
+            return true;
+        }
+        outputs.assign(1, MatShape());
+        internals.clear();
+        return true;
+    }
+
+    void getTypes(const std::vector<MatType>& inputs,
+                  const int requiredOutputs,
+                  const int requiredInternals,
+                  std::vector<MatType>& outputs,
+                  std::vector<MatType>& internals) const CV_OVERRIDE
+    {
+        CV_Assert(inputs.size() == 3);
+        int outType = (inputs[2] >= 0) ? inputs[2] : CV_32F;
+        outputs.assign(requiredOutputs, outType);
+        internals.assign(requiredInternals, MatType(-1));
+    }
+
+private:
+    template<typename Tout>
+    void runKernel(const Mat& indices, int depth, const Mat& values, Mat& Y)
+    {
+        CV_CheckGT(depth, 0, "OneHot: depth must be > 0");
+
+        MatShape outShape = shape(indices);
+        int ndims = (int)outShape.size();
+        int insAxis = normalize_axis(axis, ndims + 1);
+        outShape.insert(outShape.begin() + insAxis, depth);
+
+        if (Y.empty() || Y.type() != DataType<Tout>::type || Y.shape() != outShape)
+            Y.create(outShape, DataType<Tout>::type);
+
+        Tout offVal = (Tout)0;
+        Tout onVal  = (Tout)1;
+        if (!values.empty())
+        {
+            CV_Assert(values.total() == 2);
+            Mat flat = values.reshape(1, (int)values.total());
+            offVal = tensorToScalar<Tout>(flat.row(0));
+            onVal  = tensorToScalar<Tout>(flat.row(1));
+        }
+
+        Y.setTo(Scalar::all(offVal));
+
+        Mat idx64;
+        if (indices.depth() == CV_64S) idx64 = indices;
+        else indices.convertTo(idx64, CV_64S);
+        const int64_t* idxPtr64 = idx64.ptr<int64_t>();
+        size_t inTotal = indices.total();
+        size_t outStep = 1;
+        for (int k = insAxis + 1; k < (int)outShape.size(); ++k)
+            outStep *= (size_t)outShape[k];
+
+        Tout* outPtr = Y.ptr<Tout>();
+
+        parallel_for_(Range(0, (int)inTotal), [&](const Range& r){
+            const int64_t* localIdxPtr64 = idxPtr64;
+            Tout* localOutPtr = outPtr;
+            const size_t localOutStep = outStep;
+            const int localDepth = depth;
+            const Tout localOnVal = onVal;
+            for (int pos = r.start; pos < r.end; ++pos) {
+                int64_t idxVal = localIdxPtr64[pos];
+                int64_t wrapped = idxVal % localDepth;
+                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");
+        }
+
+        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;
+            case CV_16S:  runKernel<int16_t>(indices, (int)depth64, values, Y); break;
+            case CV_32S:  runKernel<int>(indices, (int)depth64, values, Y); break;
+            case CV_64S:  runKernel<int64_t>(indices, (int)depth64, values, Y); break;
+            case CV_16F:  runKernel<hfloat>(indices, (int)depth64, values, Y); break;
+            case CV_16BF: runKernel<bfloat>(indices, (int)depth64, values, Y); break;
+            case CV_32F:  runKernel<float>(indices, (int)depth64, values, Y); break;
+            case CV_64F:  runKernel<double>(indices, (int)depth64, values, Y); break;
+            default: CV_Error(Error::BadDepth, "OneHot: unsupported output type");
+        }
+
+        if (outKind == _InputArray::STD_VECTOR_UMAT)
+        {
+            Y.copyTo(out_umats->at(0));
+        }
+    }
+};
+
+Ptr<OneHotLayer> OneHotLayer::create(const LayerParams& params)
+{
+    return makePtr<OneHotLayerImpl>(params);
+}
+
+}} // namespace cv::dnn
@@ -217,6 +217,7 @@ protected:
    void parseTrilu                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
    void parseIsNaN                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
    void parseIsInf                (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
+    void parseOneHot               (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
    void parseDet                  (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
    void parseCenterCropPad        (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
    void parseGridSample           (LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto);
@@ -1756,6 +1757,12 @@ void ONNXImporter2::parseIsInf(LayerParams& layerParams, const opencv_onnx::Node
    addLayer(layerParams, node_proto);
 }

+void ONNXImporter2::parseOneHot(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
+{
+    layerParams.type = "OneHot";
+    addLayer(layerParams, node_proto);
+}
+
 void ONNXImporter2::parseDet(LayerParams& layerParams, const opencv_onnx::NodeProto& node_proto)
 {
    layerParams.type = "Det";
@@ -2630,6 +2637,7 @@ void ONNXImporter2::buildDispatchMap_ONNX_AI(int opset_version)
    dispatch["IsNaN"] = &ONNXImporter2::parseIsNaN;
    dispatch["IsInf"] = &ONNXImporter2::parseIsInf;
    dispatch["CenterCropPad"] = &ONNXImporter2::parseCenterCropPad;
+    dispatch["OneHot"] = &ONNXImporter2::parseOneHot;
    dispatch["Det"] = &ONNXImporter2::parseDet;
    dispatch["GridSample"] = &ONNXImporter2::parseGridSample;
    dispatch["Upsample"] = &ONNXImporter2::parseUpsample;
@@ -1666,13 +1666,13 @@ CASE(test_not_3d)
 CASE(test_not_4d)
    // no filter
 CASE(test_onehot_negative_indices)
-    // no filter
+    SKIP;
 CASE(test_onehot_with_axis)
-    // no filter
+    SKIP;
 CASE(test_onehot_with_negative_axis)
-    // no filter
+    SKIP;
 CASE(test_onehot_without_axis)
-    // no filter
+    SKIP;
 CASE(test_optional_get_element)
    // no filter
 CASE(test_optional_get_element_sequence)
@@ -572,3 +572,7 @@
 "test_gridsample_volumetric_bilinear_align_corners_1",
 "test_gridsample_volumetric_nearest_align_corners_0",
 "test_gridsample_volumetric_nearest_align_corners_1",
+"test_onehot_negative_indices",
+"test_onehot_with_axis",
+"test_onehot_with_negative_axis",
+"test_onehot_without_axis",
@@ -467,10 +467,6 @@
 "test_nesterov_momentum",  // Issues::Layer does not exist (NesterovsAcceleratedGradient) Can't create layer "onnx_node_output_0!X_new" of type "ai.onnx.preview.training.Momentum" in function 'getLayerInstance'
 "test_nllloss_NCd1d2_reduction_sum_expanded",
 "test_nllloss_NCd1d2d3d4d5_mean_weight_expanded",
-"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 ---
-"test_onehot_without_axis",  // ---- same as above ---
 "test_optional_get_element",  // Issue::out of memory :: Failed to allocate 1044051907127083008 bytes in function 'OutOfMemoryError'
 "test_optional_get_element_optional_sequence",
 "test_optional_get_element_optional_tensor",