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

Code before:
                            int D = -1, H = -1, W = -1;
                            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 = sz.at<int32_t>(2); H = sz.at<int32_t>(3); W = sz.at<int32_t>(4); }
                            if (D > 0 && H > 0 && W > 0) resolved = MatShape{N, D, H, W, 3};
                        }
                    }
                } catch (const cv::Exception& e) {
                    CV_Error(cv::Error::StsError, cv::format("DNN/AffineGrid: failed to resolve output shape from 'size' tensor: %s", e.what()));
                }
            }
            outShape = resolved;
        } else {
            if (is3d)
                outShape = MatShape{N, -1, -1, -1, 3};
            else
                outShape = MatShape{N, -1, -1, 2};
        }
        outputs.assign(1, outShape);
        return false;
    }

    bool dynamicOutputShapes() const CV_OVERRIDE
    {
        Net::Impl* netimpl_ = getNetImpl(this);
        if (!netimpl_) return true;
        if (this->inputs.size() < 2) return true;

        if (netimpl_->isConstArg(this->inputs[1])) return false;
        try {
            const Mat& sz = netimpl_->argTensor(this->inputs[1]);
            return sz.empty();
        } catch (const cv::Exception& e) {
            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(), 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, "");


Review comment (suggestion):
affineGrid requires two inputs, otherwise an error should be reported.

Code after:

template<typename T>
void computeGrid3D(const Mat& theta, Mat& grid, int N, int D, int H, int W,
                   float xs, float xd, float ys, float yd, float zs, float zd)
{
    const int stride = D * H;
    const int total = N * stride;
    parallel_for_(Range(0, total), [&](const Range& range){
        const int H_ = H;
        const int W_ = W;
        const int stride_ = stride;
        const float xs_ = xs, xd_ = xd, ys_ = ys, yd_ = yd, zs_ = zs, zd_ = zd;
        const Mat theta_ = theta;
        Mat grid_ = grid;
        int lastN = -1;
        const T* T0 = nullptr;
        const T* T1 = nullptr;
        const T* T2 = nullptr;
        float T0_0 = 0.f, T1_0 = 0.f, T2_0 = 0.f;
        for (int ndz = range.start; ndz < range.end; ndz++)
        {
            int n = ndz / stride_;
            int rem = ndz - n * stride_;
            int z = rem / H_;
            int y = rem - z * H_;
            if (n != lastN)
            {
                int i0[3] = {n, 0, 0};
                int i1[3] = {n, 1, 0};
                int i2[3] = {n, 2, 0};
                T0 = theta_.ptr<T>(i0);
                T1 = theta_.ptr<T>(i1);
                T2 = theta_.ptr<T>(i2);
                T0_0 = (float)T0[0];
                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;