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

Code before:

    bool supportBackend(int backendId) CV_OVERRIDE
    {
        return backendId == DNN_BACKEND_OPENCV;
    }

    bool dynamicOutputShapes() const CV_OVERRIDE
    {
        return false;
    }

    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() >= 2);
        const MatShape& idxShape = inputs[0];
        CV_Assert(!idxShape.empty());

        try
        {
            if ((int)this->inputs.size() >= 2)
            {
                Net::Impl* netimpl_ = getNetImpl(this);
                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)
                    {
                        MatShape outShape = idxShape;
                        int insAxis = normalize_axis(axis, (int)outShape.size() + 1);
                        outShape.insert(outShape.begin() + insAxis, (int)depth64);
                        outputs.assign(1, outShape);
                        internals.clear();
                        return true;
                    }
                }
            }
        }
        catch (const cv::Exception& e)
        {
            CV_LOG_DEBUG(NULL, "OneHot: getMemoryShapes dynamic fallback (cv::Exception): " << e.what());
        }

        outputs.assign(1, MatShape());
        internals.clear();
        return true;
    }

Review comment (refactor):
I suggest to extract this logic into a dedicated method that would return the shape of output tensor based on the shape of input and the 'depth' tensor.

Code after:

        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());