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(), 1, "");
        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");

        if (inputs.size() >= 2)
            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
    {
        std::vector<Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

        const Mat& theta = inputs[0];
        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;
        if ((int)inputs.size() >= 2) {
            Mat sz = inputs[1];
            CV_CheckTrue(sz.total() == (size_t)(is3d ? 5 : 4), "size input must have 4 (2D) or 5 (3D) elements");
            if (is3d) {
                if (sz.depth() == CV_64S) {
                    D = (int)sz.at<int64_t>(2);
                    H = (int)sz.at<int64_t>(3);
                    W = (int)sz.at<int64_t>(4);
                } else if (sz.depth() == CV_32S) {
                    D = (int)sz.at<int32_t>(2);
                    H = (int)sz.at<int32_t>(3);
                    W = (int)sz.at<int32_t>(4);
                } else {
                    CV_CheckType(sz.depth(), sz.depth() == CV_64S || sz.depth() == CV_32S, "");
                }
            } else {
                if (sz.depth() == CV_64S) {

Review comment (suggestion):
because the second input is not necessarily a constant tensor, in this case outputs_arr must be properly allocated. So, don't use getMatVector. Instead, check kind of inputs_arr/outputs_arr and add two branches, depending on whether inputs_arr/output_arr are STD_VECTOR_MAT's or STD_VECTOR_UMAT's.

Code after:
                T1_0 = (float)T1[0];
                T2_0 = (float)T2[0];
                lastN = n;
            }
            float ny = yd_ + ys_ * y;
            float nz = zd_ + zs_ * z;
            float base0 = (float)T0[1]*ny + (float)T0[2]*nz + (float)T0[3];
            float base1 = (float)T1[1]*ny + (float)T1[2]*nz + (float)T1[3];
            float base2 = (float)T2[1]*ny + (float)T2[2]*nz + (float)T2[3];
            int idx[5] = {n, z, y, 0, 0};
            float* out = grid_.ptr<float>(idx);
            for (int x = 0; x < W_; x++)
            {
                float nx = xd_ + xs_ * x;
                int o = 3*x;
                out[o + 0] = T0_0*nx + base0;
                out[o + 1] = T1_0*nx + base1;
                out[o + 2] = T2_0*nx + base2;
            }
        }
    });
}
}

class AffineGridLayerImpl CV_FINAL : public AffineGridLayer
{
public:
    AffineGridLayerImpl(const LayerParams& params)
    {
        setParamsFrom(params);
        align_corners = params.get<bool>("align_corners", false);
    }

private:
    void resolveSizeAndShape(int N, bool is3d, const Mat* explicitSz,
                         MatShape& outShape, int& D, int& H, int& W) const
    {
        outShape = is3d ? MatShape{N, -1, -1, -1, 3} : MatShape{N, -1, -1, 2};
        D = 1; H = -1; W = -1;

        Mat sz;
        if (explicitSz)
        {
            sz = *explicitSz;
        }
        else
        {
            Net::Impl* netimpl_ = getNetImpl(this);
            if (!netimpl_)
                return;
            try