d2n_estimator.cpp

/* ***********************************************************
 *
 * Description:
 *  d2n_estimator.h/cpp estimates surface normals from a depth map using
 *  classical image processing (no ML).
 *
 * Procedure (3 steps):
 *  1) Depth -> XYZ (range image) + valid mask   [pinhole back-projection]
 *  2) XYZ   -> dU, dV (tangent vectors)        [derivative filters]
 *  3) dU,dV -> normals                         [cross product + normalize]
 *
 * Notes:
 *  - Dataset convention: Z == 1.0 denotes background (invalid).
 *  - A "safe" mask is built (erosion) so derivatives are not computed across
 *    valid/invalid boundaries.
 *    
 * ----------------------------------------------------------
 *
 *  Copyright (C) 2026 Claudio Z. (cloudofoz)
 *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy
 *  of this software and associated documentation files (the "Software"), to deal
 *  in the Software without restriction, including without limitation the rights
 *  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 *  copies of the Software, and to permit persons to whom the Software is
 *  furnished to do so, subject to the following conditions:
 *
 *  The above copyright notice and this permission notice shall be included in all
 *  copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 *  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 *  SOFTWARE.
 *
 * **********************************************************/

#include "d2n_estimator.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cfloat>
#include "opencv2/opencv.hpp"

 //---------------------------------------------------------
 // Dataset-specific validity

static inline bool is_valid_depth(float Z)
{
    // Dataset-specific convention:
    // Z == 1.0 denotes background pixels, Z > 1.0 is valid geometry.
    return Z > 1.0f;
}

//---------------------------------------------------------
// Step 1: Depth -> XYZ + validMask
//
// This function converts a depth image into a 3D point map using a pinhole camera model.
// Each pixel with a valid depth value is back-projected into camera space (X, Y, Z).
// At the same time, a binary mask is produced to explicitly track which pixels
// correspond to valid 3D points.
int d2n_depth_to_xyz(const cv::Mat& depth,
    const CameraIntrinsics& K,
    cv::Mat& outXYZ,
    cv::Mat& validMask)
{
    // The function expects a single-channel floating-point depth map.
    // Each pixel encodes the distance along the camera Z axis.
    CV_Assert(depth.type() == CV_32FC1);

    // Allocate output buffers:
    // - outXYZ stores a 3D point (X,Y,Z) for each pixel
    // - validMask marks which pixels correspond to a valid depth measurement
    outXYZ.create(depth.rows, depth.cols, CV_32FC3);
    validMask.create(depth.rows, depth.cols, CV_8UC1);

    int validPixels = 0;

    // Iterate over the image grid.
    // Pixel coordinates (u, v) are interpreted as image-space coordinates.
    for (int v = 0; v < depth.rows; ++v)
    {
        for (int u = 0; u < depth.cols; ++u)
        {
            const float Z = depth.at<float>(v, u);

            // Invalid depth values (NaNs, infinities, or values outside
            // the accepted range) are explicitly discarded.
            // This avoids propagating meaningless geometry further in the pipeline.
            if (!std::isfinite(Z) || !is_valid_depth(Z))
            {
                outXYZ.at<cv::Vec3f>(v, u) = cv::Vec3f(0, 0, 0);
                validMask.at<uchar>(v, u) = 0;
                continue;
            }

            // Pinhole back-projection.
            // The pixel is reinterpreted as a ray passing through the camera center,
            // and scaled by the measured depth Z.
            //
            // This converts image coordinates (u, v, Z) into metric 3D coordinates
            // expressed in the camera reference frame.
            // X = (u - cx) * Z / fx
            // Y = (v - cy) * Z / fy
            const float X = (static_cast<float>(u) - K.cx) * Z / K.fx;
            const float Y = (static_cast<float>(v) - K.cy) * Z / K.fy;

            // Store the reconstructed 3D point and mark the pixel as valid.
            outXYZ.at<cv::Vec3f>(v, u) = cv::Vec3f(X, Y, Z);
            validMask.at<uchar>(v, u) = 255;
            ++validPixels;
        }
    }

    // Return the total number of valid 3D points.
    // This value is often useful for later normalization or statistics.
    return validPixels;
}

//---------------------------------------------------------
// Step 2: XYZ -> derivatives (dU, dV)
//
// Goal: estimate how the 3D surface changes when moving by one pixel in the image.
// In practice we compute two vector fields:
//   - out_dU: change of the 3D point when moving horizontally (u direction)
//   - out_dV: change of the 3D point when moving vertically   (v direction)
//
// Each derivative is a 3D vector (dX, dY, dZ). These are the basic data
// for building surface normals in the next step.
void d2n_xyz_to_derivatives(const cv::Mat& xyz,
    const cv::Mat& validMask,
    cv::Mat& out_dU,
    cv::Mat& out_dV,
    const D2NEstimatorParams& params)
{
    // Basic sanity checks
    CV_Assert(!xyz.empty());
    CV_Assert(xyz.type() == CV_32FC3);
    CV_Assert(validMask.type() == CV_8UC1);
    CV_Assert(xyz.size() == validMask.size());

    // 1) Split the 3D point map into three scalar images (X, Y, Z).
    // Doing this makes it easy to apply standard 2D derivative operators channel-wise.
    std::array<cv::Mat, 3> xyzCh;
    cv::split(xyz, xyzCh.data()); // each channel is CV_32FC1

    // Temporary storage for partial derivatives of each channel.
    // dU[i] and dV[i] will store the derivative of the i-th coordinate (X,Y,Z).
    std::array<cv::Mat, 3> dU, dV;

    // 2) Apply a discrete derivative operator.
    // Different operators behave slightly differently:
    // but they all estimate "how much does this value change across pixels".
    //
    // supportRadius how far (in pixels) the operator "reaches" from the center.
    // Used later to shrink the valid mask, discard the derivatives that mix
    // valid and invalid measurements.
    int supportRadius = 1; // default for fixed 3x3 operators

    switch (params.filter)
    {
    case D2NEstimatorParams::DerivativeFilter::Scharr:
    {
        // Scharr in OpenCV is a fixed 3x3 operator 
        // (often a bit more stable than Sobel at 3x3).
        CV_Assert(params.ksize == 3);
        supportRadius = 1;

        for (int i = 0; i < 3; ++i)
        {
            // d/du: horizontal derivative (x-direction in image space)
            cv::Scharr(xyzCh[i], dU[i], CV_32F, 1, 0);
            // d/dv: vertical derivative (y-direction in image space)
            cv::Scharr(xyzCh[i], dV[i], CV_32F, 0, 1);
        }
    } break;

    case D2NEstimatorParams::DerivativeFilter::Prewitt:
    {
        // Prewitt is implemented here explicitly with fixed 3x3 kernels.
        // It is simple and often used as a baseline derivative operator.
        CV_Assert(params.ksize == 3);
        supportRadius = 1;

        static const cv::Mat kx = (cv::Mat_<float>(3, 3) <<
            -1, 0, 1,
            -1, 0, 1,
            -1, 0, 1);

        static const cv::Mat ky = (cv::Mat_<float>(3, 3) <<
            -1, -1, -1,
            0, 0, 0,
            1, 1, 1);

        for (int i = 0; i < 3; ++i)
        {
            // filter2D applies the kernel as-is and returns the discrete derivative estimate.
            cv::filter2D(xyzCh[i], dU[i], CV_32F, kx);
            cv::filter2D(xyzCh[i], dV[i], CV_32F, ky);
        }
    } break;

    case D2NEstimatorParams::DerivativeFilter::Sobel:
    default:
    {
        // Sobel is the standard choice and supports larger odd kernel sizes.
        // Larger kernels tend to be smoother (less noisy) but also blur fine detail.
        CV_Assert(params.ksize == 1 || params.ksize == 3 || params.ksize == 5 || params.ksize == 7);
        supportRadius = (params.ksize - 1) / 2;

        for (int i = 0; i < 3; ++i)
        {
            cv::Sobel(xyzCh[i], dU[i], CV_32F, 1, 0, params.ksize); // horizontal derivative
            cv::Sobel(xyzCh[i], dV[i], CV_32F, 0, 1, params.ksize); // vertical derivative
        }
    } break;
    }

    // 3) Build a "safe" mask: derivatives near invalid pixels are unreliable.
    // A derivative operator compares a pixel with its neighbors; if any neighbor is invalid,
    // the derivative becomes contaminated. Eroding the mask removes a border around invalid
    // regions, keeping only pixels whose whole filter neighborhood is valid.
    cv::Mat safeMask = validMask;

    int maskErodeIters = params.maskErodeIters;
    if (maskErodeIters < 0)
        maskErodeIters = supportRadius; // auto: match the filter footprint

    // Not erode more than the operator radius (anything more would be unnecessarily strict).
    maskErodeIters = std::min(maskErodeIters, supportRadius);

    if (maskErodeIters > 0)
        cv::erode(validMask, safeMask, cv::Mat(), cv::Point(-1, -1), maskErodeIters);

    // 4) Merge the per-channel derivatives back into vector fields (dX,dY,dZ).
    // After this, out_dU and out_dV are CV_32FC3 images aligned with the input.
    cv::merge(dU.data(), 3, out_dU); // CV_32FC3
    cv::merge(dV.data(), 3, out_dV); // CV_32FC3

    // 5) Explicitly zero derivatives outside the safe region.
    // This keeps next computations (e.g., normals) from using derivatives that we
    // already know are unreliable, and makes invalid areas visually and numerically clean.
    out_dU.setTo(cv::Scalar(0, 0, 0), safeMask == 0);
    out_dV.setTo(cv::Scalar(0, 0, 0), safeMask == 0);
}


//---------------------------------------------------------
// Step 3: derivatives -> normals
//
// Goal: convert the two local surface "tangent directions" (dU and dV) into a unit normal map.
//
// - dU tells how the 3D surface point changes when moving 1 pixel along u (image x direction).
// - dV tells how the 3D surface point changes when moving 1 pixel along v (image y direction).
// These two vectors lie on (or approximate) the local tangent plane of the surface.
// The normal is the vector perpendicular to that plane, obtained via the cross product.
void d2n_derivatives_to_normals(const cv::Mat& dU,
    const cv::Mat& dV,
    const cv::Mat& validMask,
    cv::Mat& outNormals,
    bool enforcePositiveZ /*= true*/)
{
    // Sanity checks
    CV_Assert(dU.type() == CV_32FC3);
    CV_Assert(dV.type() == CV_32FC3);
    CV_Assert(validMask.type() == CV_8UC1);
    CV_Assert(dU.size() == dV.size());
    CV_Assert(dU.size() == validMask.size());

    // Output is a per-pixel unit vector (nx, ny, nz).
    outNormals.create(dU.rows, dU.cols, CV_32FC3);

    for (int v = 0; v < dU.rows; ++v)
    {
        // Use row pointers for efficiency: this is a tight per-pixel loop.
        const cv::Vec3f* dUPtr = dU.ptr<cv::Vec3f>(v);
        const cv::Vec3f* dVPtr = dV.ptr<cv::Vec3f>(v);
        const uchar* maskPtr = validMask.ptr<uchar>(v);
        cv::Vec3f* nPtr = outNormals.ptr<cv::Vec3f>(v);

        for (int u = 0; u < dU.cols; ++u)
        {
            // If the pixel is invalid, keep the output clean.
            if (maskPtr[u] == 0)
            {
                nPtr[u] = cv::Vec3f(0, 0, 0);
                continue;
            }

            // Compute the normal as the vector orthogonal to the local tangents:
            // n_raw = dU x dV
            //
            // Note: the order matters (dU x dV vs dV x dU), it flips the direction.
            cv::Vec3f n = dUPtr[u].cross(dVPtr[u]);

            // Normalize to obtain a unit vector.
            // If the magnitude is near zero, the direction is unstable (e.g. invalid geometry)
            const float length = std::sqrt(n.dot(n));
            if (length <= FLT_EPSILON)
            {
                nPtr[u] = cv::Vec3f(0, 0, 0);
                continue; // zero-length vector -> undefined normal
            }

            n /= length;

            // Optional orientation convention:
            // Some datasets/benchmarks assume normals face the camera (positive Z in camera space),
            if (enforcePositiveZ && n[2] < 0.0f)
                n = -n;

            nPtr[u] = n;
        }
    }
}


//---------------------------------------------------------

// High-level normal estimation function
//
// This function connects all the individual steps
// of the depth-to-normals pipeline.
//
// Data Flow:
//
//   depth image
//        -> 3D point map (XYZ)
//        -> spatial derivatives (dU, dV)
//        -> surface normals
//
// Each step operates on the output of the previous one.
bool d2n_estimate(const Dataset& ds,
    const DatasetFrame& inFrame,
    const D2NEstimatorParams& params,
    D2NEstimatorOutput& outResult)
{
    // Camera intrinsics are taken from the dataset.
    // They define the geometric relationship between pixel coordinates
    // and metric 3D space.
    const CameraIntrinsics& K = ds.params.K;

    // Step 1: Depth -> XYZ
    //
    // Convert the depth map into a 3D point cloud expressed in the camera frame.
    // At the same time, build a validity mask and count how many pixels produce
    // a meaningful 3D measurement.
    outResult.valid_pixels =
        d2n_depth_to_xyz(inFrame.depth,
            K,
            outResult.xyz,
            outResult.valid_mask);

    // Step 2: XYZ -> derivatives
    //
    // Estimate how the 3D surface changes across neighboring pixels.
    // The output consists of two vector fields (du, dv) that approximate
    // the local tangent directions of the surface.
    d2n_xyz_to_derivatives(outResult.xyz,
        outResult.valid_mask,
        outResult.du,
        outResult.dv,
        params);

    // Step 3: derivatives -> normals
    //
    // Convert local tangent directions into unit surface normals
    // using the cross product.
    d2n_derivatives_to_normals(outResult.du,
        outResult.dv,
        outResult.valid_mask,
        outResult.estimated_normals,
        /*enforcePositiveZ=*/false);

    return true;
}