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

Code before:
    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_CheckGE((int)inputs.size(), 2, "AffineGrid requires two inputs: theta and size");
        const bool allow_fp16 = (preferableTarget == DNN_TARGET_OPENCL_FP16 || preferableTarget == DNN_TARGET_CPU_FP16);
        const int thetaType = inputs[0];
        const bool is_float_ok = thetaType == CV_32F || thetaType == CV_64F || thetaType == CV_16BF || (allow_fp16 && thetaType == CV_16F) || thetaType == CV_16F;
        CV_CheckType(thetaType, is_float_ok, "AffineGrid: theta must be a floating tensor");
        CV_CheckType(inputs[1], inputs[1] == CV_64S || inputs[1] == CV_32S, "");

        outputs.assign(std::max(1, requiredOutputs), CV_32F);
        internals.assign(requiredInternals, CV_32F);
    }

    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        auto inKind = inputs_arr.kind();
        auto outKind = outputs_arr.kind();
        CV_Assert(inKind == _InputArray::STD_VECTOR_MAT || inKind == _InputArray::STD_VECTOR_UMAT);
        CV_Assert(outKind == _InputArray::STD_VECTOR_MAT || outKind == _InputArray::STD_VECTOR_UMAT);

        Mat theta, sz;
        if (inKind == _InputArray::STD_VECTOR_MAT) {
            theta = inputs_arr.getMat(0);
            sz = inputs_arr.getMat(1);
        } else {
            UMat utheta = inputs_arr.getUMat(0);
            UMat usz = inputs_arr.getUMat(1);
            theta = utheta.getMat(ACCESS_READ);
            sz = usz.getMat(ACCESS_READ);
        }

        CV_CheckTrue(theta.dims == 3, "theta must be 3D [N, r, c]");
        const int N = theta.size[0];
        const int r = theta.size[1];
        const bool is3d = (r == 3);

        int D = 1, H = 0, W = 0;
        MatShape dummy;
        resolveSizeAndShape(N, is3d, &sz, dummy, D, H, W);

        CV_CheckGT(H, 0, "H must be > 0");
        CV_CheckGT(W, 0, "W must be > 0");
        if (is3d) CV_CheckGT(D, 0, "D must be > 0");

        std::vector<int> outShape = is3d ? std::vector<int>{N, D, H, W, 3}
                                         : std::vector<int>{N, H, W, 2};
        MatShape fitShape(outShape.begin(), outShape.end());

Review comment (suggestion):
for inputs_arr it's enough to use getMatVector. No need to have this branch. Moreover, this code with getMatVector will also be compatible with GpuMat and any other kind of input arrays. This dispatcher is necessary only for the outputs.

Code after:
            CV_Error(cv::Error::StsError, cv::format("DNN/AffineGrid: cannot query 'size' tensor at runtime: %s", e.what()));
        }
    }

    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_CheckGE((int)inputs.size(), 2, "AffineGrid requires two inputs: theta and size");
        const bool allow_fp16 = (preferableTarget == DNN_TARGET_OPENCL_FP16 || preferableTarget == DNN_TARGET_CPU_FP16);
        const int thetaType = inputs[0];
        const bool is_float_ok = thetaType == CV_32F || thetaType == CV_64F || thetaType == CV_16BF || (allow_fp16 && thetaType == CV_16F) || thetaType == CV_16F;
        CV_CheckType(thetaType, is_float_ok, "AffineGrid: theta must be a floating tensor");
        CV_CheckType(inputs[1], inputs[1] == CV_64S || inputs[1] == CV_32S, "");

        outputs.assign(std::max(1, requiredOutputs), CV_32F);
        internals.assign(requiredInternals, CV_32F);
    }

    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        auto outKind = outputs_arr.kind();
        CV_Assert(outKind == _InputArray::STD_VECTOR_MAT || outKind == _InputArray::STD_VECTOR_UMAT);

        std::vector<Mat> inMats;
        inputs_arr.getMatVector(inMats);
        CV_CheckGE((int)inMats.size(), 2, "AffineGrid requires two inputs: theta and size");
        Mat theta = inMats[0];
        Mat sz = inMats[1];

        CV_CheckTrue(theta.dims == 3, "theta must be 3D [N, r, c]");
        const int N = theta.size[0];
        const int r = theta.size[1];
        const bool is3d = (r == 3);

        int D = 1, H = 0, W = 0;
        MatShape dummy;
        resolveSizeAndShape(N, is3d, &sz, dummy, D, H, W);

        CV_CheckGT(H, 0, "H must be > 0");
        CV_CheckGT(W, 0, "W must be > 0");
        if (is3d) CV_CheckGT(D, 0, "D must be > 0");

        std::vector<int> outShape = is3d ? std::vector<int>{N, D, H, W, 3}
                                         : std::vector<int>{N, H, W, 2};
        MatShape fitShape(outShape.begin(), outShape.end());
        Mat grid;
        if (outKind == _InputArray::STD_VECTOR_MAT) {
            std::vector<Mat>& outs = outputs_arr.getMatVecRef();