Add OpenCV pinhole camera rays (#3901)

Author: liorbenhorin

changelog/+opencv-pinhole-rays-7c2f4a91.added.md

@@ -0,0 +1 @@
+Add `SensorTiledCamera.utils.compute_camera_rays_pinhole_opencv()` for OpenCV pinhole cameras with radial, tangential, and thin-prism distortion.

docs/concepts/sensors.rst

@@ -117,8 +117,8 @@ attributes, and usage examples.
 * :class:`~newton.sensors.SensorIMU` -- linear acceleration and angular velocity at site frames.
 * :class:`~newton.sensors.SensorTiledCamera` -- raytraced color and depth rendering across multiple worlds.
 
-Camera Rays from USD Data
--------------------------
+Camera Rays from USD and Calibration Data
+-----------------------------------------
 
 ``SensorTiledCamera`` can build standard USD pinhole camera rays directly. For lens models without standard USD
 attributes, read the attributes you use in your pipeline and pass the numeric values into the matching helper:
@@ -144,6 +144,10 @@ attributes, read the attributes you use in your pipeline and pass the numeric va
        color_image=color,
    )
 
+For OpenCV-calibrated pinhole cameras, call
+:meth:`~newton.sensors.SensorTiledCamera.Utils.compute_camera_rays_pinhole_opencv` with the calibrated intrinsics and
+radial, tangential, and optional thin-prism coefficients.
+
 For fisheye cameras, extract the calibration values from your chosen USD attributes and call one of
 :meth:`~newton.sensors.SensorTiledCamera.Utils.compute_camera_rays_fisheye_opencv`,
 :meth:`~newton.sensors.SensorTiledCamera.Utils.compute_camera_rays_fisheye_ftheta`, or

newton/_src/sensors/warp_raytrace/camera_utils.py

@@ -440,6 +440,232 @@ def compute_camera_rays_pinhole_from_aperture_kernel(
     out_rays[output_camera_index, py, px, 1] = wp.normalize(ray_direction_camera_space)
 
 
+_PINHOLE_OPENCV_INVERSION_ITERATIONS = 10
+_PINHOLE_OPENCV_LINE_SEARCH_ITERATIONS = 6
+# Residuals are in normalized image coordinates; convergence ~1e-6, acceptance ~2e-5 RMS.
+_PINHOLE_OPENCV_CONVERGENCE_TOLERANCE_SQUARED = 1.0e-12
+_PINHOLE_OPENCV_ACCEPTANCE_TOLERANCE_SQUARED = 4.0e-10
+_PINHOLE_OPENCV_SOLVER_EPSILON = 1.0e-8
+
+
+@wp.func
+def _distort_pinhole_opencv(
+    x: wp.float32,
+    y: wp.float32,
+    k1: wp.float32,
+    k2: wp.float32,
+    k3: wp.float32,
+    k4: wp.float32,
+    k5: wp.float32,
+    k6: wp.float32,
+    p1: wp.float32,
+    p2: wp.float32,
+    s1: wp.float32,
+    s2: wp.float32,
+    s3: wp.float32,
+    s4: wp.float32,
+) -> wp.vec2f:
+    radius_squared = x * x + y * y
+    radius_fourth = radius_squared * radius_squared
+    radius_sixth = radius_fourth * radius_squared
+    radial = (1.0 + k1 * radius_squared + k2 * radius_fourth + k3 * radius_sixth) / (
+        1.0 + k4 * radius_squared + k5 * radius_fourth + k6 * radius_sixth
+    )
+    tangential_prism_x = (
+        2.0 * p1 * x * y + p2 * (radius_squared + 2.0 * x * x) + s1 * radius_squared + s2 * radius_fourth
+    )
+    tangential_prism_y = (
+        p1 * (radius_squared + 2.0 * y * y) + 2.0 * p2 * x * y + s3 * radius_squared + s4 * radius_fourth
+    )
+    return wp.vec2f(x * radial + tangential_prism_x, y * radial + tangential_prism_y)
+
+
+@wp.func
+def _invert_pinhole_opencv(
+    x_distorted: wp.float32,
+    y_distorted: wp.float32,
+    k1: wp.float32,
+    k2: wp.float32,
+    k3: wp.float32,
+    k4: wp.float32,
+    k5: wp.float32,
+    k6: wp.float32,
+    p1: wp.float32,
+    p2: wp.float32,
+    s1: wp.float32,
+    s2: wp.float32,
+    s3: wp.float32,
+    s4: wp.float32,
+) -> wp.vec3f:
+    # Damped Newton inversion; unverifiable pixels return the zero sentinel.
+    x = x_distorted
+    y = y_distorted
+    valid = True
+    done = False
+    for _ in range(_PINHOLE_OPENCV_INVERSION_ITERATIONS):
+        if valid and not done:
+            distorted = _distort_pinhole_opencv(x, y, k1, k2, k3, k4, k5, k6, p1, p2, s1, s2, s3, s4)
+            residual_x = distorted[0] - x_distorted
+            residual_y = distorted[1] - y_distorted
+            residual_squared = residual_x * residual_x + residual_y * residual_y
+
+            if not wp.isfinite(residual_squared):
+                valid = False
+            elif residual_squared <= _PINHOLE_OPENCV_CONVERGENCE_TOLERANCE_SQUARED:
+                done = True
+            else:
+                radius_squared = x * x + y * y
+                radius_fourth = radius_squared * radius_squared
+                radius_sixth = radius_fourth * radius_squared
+                numerator = 1.0 + k1 * radius_squared + k2 * radius_fourth + k3 * radius_sixth
+                denominator = 1.0 + k4 * radius_squared + k5 * radius_fourth + k6 * radius_sixth
+
+                if wp.abs(denominator) <= _PINHOLE_OPENCV_SOLVER_EPSILON:
+                    valid = False
+                else:
+                    radial = numerator / denominator
+                    numerator_derivative = k1 + 2.0 * k2 * radius_squared + 3.0 * k3 * radius_fourth
+                    denominator_derivative = k4 + 2.0 * k5 * radius_squared + 3.0 * k6 * radius_fourth
+                    radial_derivative = (numerator_derivative * denominator - numerator * denominator_derivative) / (
+                        denominator * denominator
+                    )
+                    radial_derivative_x = 2.0 * x * radial_derivative
+                    radial_derivative_y = 2.0 * y * radial_derivative
+
+                    jacobian_00 = (
+                        radial
+                        + x * radial_derivative_x
+                        + 2.0 * p1 * y
+                        + 6.0 * p2 * x
+                        + 2.0 * s1 * x
+                        + 4.0 * s2 * radius_squared * x
+                    )
+                    jacobian_01 = (
+                        x * radial_derivative_y
+                        + 2.0 * p1 * x
+                        + 2.0 * p2 * y
+                        + 2.0 * s1 * y
+                        + 4.0 * s2 * radius_squared * y
+                    )
+                    jacobian_10 = (
+                        y * radial_derivative_x
+                        + 2.0 * p1 * x
+                        + 2.0 * p2 * y
+                        + 2.0 * s3 * x
+                        + 4.0 * s4 * radius_squared * x
+                    )
+                    jacobian_11 = (
+                        radial
+                        + y * radial_derivative_y
+                        + 6.0 * p1 * y
+                        + 2.0 * p2 * x
+                        + 2.0 * s3 * y
+                        + 4.0 * s4 * radius_squared * y
+                    )
+                    determinant = jacobian_00 * jacobian_11 - jacobian_01 * jacobian_10
+
+                    if not wp.isfinite(determinant) or wp.abs(determinant) <= _PINHOLE_OPENCV_SOLVER_EPSILON:
+                        valid = False
+                    else:
+                        delta_x = (jacobian_11 * residual_x - jacobian_01 * residual_y) / determinant
+                        delta_y = (jacobian_00 * residual_y - jacobian_10 * residual_x) / determinant
+                        if not wp.isfinite(delta_x) or not wp.isfinite(delta_y):
+                            valid = False
+                        else:
+                            base_x = x
+                            base_y = y
+                            step = 1.0
+                            accepted = False
+                            for _line_search_iteration in range(_PINHOLE_OPENCV_LINE_SEARCH_ITERATIONS):
+                                candidate_x = base_x - step * delta_x
+                                candidate_y = base_y - step * delta_y
+                                candidate_distorted = _distort_pinhole_opencv(
+                                    candidate_x,
+                                    candidate_y,
+                                    k1,
+                                    k2,
+                                    k3,
+                                    k4,
+                                    k5,
+                                    k6,
+                                    p1,
+                                    p2,
+                                    s1,
+                                    s2,
+                                    s3,
+                                    s4,
+                                )
+                                candidate_residual_x = candidate_distorted[0] - x_distorted
+                                candidate_residual_y = candidate_distorted[1] - y_distorted
+                                candidate_residual_squared = (
+                                    candidate_residual_x * candidate_residual_x
+                                    + candidate_residual_y * candidate_residual_y
+                                )
+                                if (
+                                    not accepted
+                                    and wp.isfinite(candidate_residual_squared)
+                                    and candidate_residual_squared < residual_squared
+                                ):
+                                    x = candidate_x
+                                    y = candidate_y
+                                    accepted = True
+                                step *= 0.5
+                            # Stalled: no step reduced the residual; stop iterating.
+                            if not accepted:
+                                done = True
+
+    # Accept only if the forward model reproduces the target within tolerance.
+    if valid:
+        distorted = _distort_pinhole_opencv(x, y, k1, k2, k3, k4, k5, k6, p1, p2, s1, s2, s3, s4)
+        residual_x = distorted[0] - x_distorted
+        residual_y = distorted[1] - y_distorted
+        residual_squared = residual_x * residual_x + residual_y * residual_y
+        valid = wp.isfinite(residual_squared) and residual_squared <= _PINHOLE_OPENCV_ACCEPTANCE_TOLERANCE_SQUARED
+
+    if valid:
+        return wp.normalize(wp.vec3f(x, -y, -1.0))
+    return wp.vec3f(0.0)
+
+
+@wp.kernel(enable_backward=False)
+def compute_camera_rays_pinhole_opencv_kernel(
+    width: int,
+    height: int,
+    image_width: wp.float32,
+    image_height: wp.float32,
+    fx: wp.float32,
+    fy: wp.float32,
+    cx: wp.float32,
+    cy: wp.float32,
+    k1: wp.float32,
+    k2: wp.float32,
+    k3: wp.float32,
+    k4: wp.float32,
+    k5: wp.float32,
+    k6: wp.float32,
+    p1: wp.float32,
+    p2: wp.float32,
+    s1: wp.float32,
+    s2: wp.float32,
+    s3: wp.float32,
+    s4: wp.float32,
+    camera_index: int,
+    out_rays: wp.array4d[wp.vec3f],
+):
+    py, px = wp.tid()
+    u = ((float(px) + 0.5) / float(width)) * image_width
+    v = ((float(py) + 0.5) / float(height)) * image_height
+    x_distorted = (u - cx) / fx
+    y_distorted = (v - cy) / fy
+
+    ray_direction_camera_space = _invert_pinhole_opencv(
+        x_distorted, y_distorted, k1, k2, k3, k4, k5, k6, p1, p2, s1, s2, s3, s4
+    )
+
+    out_rays[camera_index, py, px, 0] = wp.vec3f(0.0)
+    out_rays[camera_index, py, px, 1] = ray_direction_camera_space
+
+
 @wp.kernel(enable_backward=False)
 def compute_camera_rays_fisheye_opencv_kernel(
     width: int,

newton/_src/sensors/warp_raytrace/utils.py

@@ -674,6 +674,124 @@ def compute_camera_rays_usd_pinhole(
             camera_index=camera_index,
         )
 
+    def compute_camera_rays_pinhole_opencv(
+        self,
+        width: int,
+        height: int,
+        fx: float,
+        fy: float,
+        cx: float,
+        cy: float,
+        *,
+        image_width: float | None = None,
+        image_height: float | None = None,
+        k1: float = 0.0,
+        k2: float = 0.0,
+        k3: float = 0.0,
+        k4: float = 0.0,
+        k5: float = 0.0,
+        k6: float = 0.0,
+        p1: float = 0.0,
+        p2: float = 0.0,
+        s1: float = 0.0,
+        s2: float = 0.0,
+        s3: float = 0.0,
+        s4: float = 0.0,
+        out_rays: wp.array4d[wp.vec3f] | None = None,
+        camera_index: int = 0,
+    ) -> wp.array4d[wp.vec3f]:
+        """Compute camera-space ray directions for OpenCV pinhole cameras.
+
+        Inverts OpenCV's rational radial, tangential, and thin-prism distortion
+        model with damped Newton iteration. The four- and five-coefficient
+        variants are represented by leaving unused coefficients at zero. The
+        forward distortion mapping must be invertible over the calibrated image
+        domain. Pixels whose inverse cannot be verified within the solver
+        tolerance receive a zero direction.
+
+        Args:
+            width: Output image width [px].
+            height: Output image height [px].
+            fx: Horizontal focal length [px].
+            fy: Vertical focal length [px].
+            cx: Principal point x-coordinate [px].
+            cy: Principal point y-coordinate [px].
+            image_width: Calibration image width [px]. If ``None``, uses
+                *width*.
+            image_height: Calibration image height [px]. If ``None``, uses
+                *height*.
+            k1: First numerator radial distortion coefficient.
+            k2: Second numerator radial distortion coefficient.
+            k3: Third numerator radial distortion coefficient.
+            k4: First denominator radial distortion coefficient.
+            k5: Second denominator radial distortion coefficient.
+            k6: Third denominator radial distortion coefficient.
+            p1: First tangential distortion coefficient.
+            p2: Second tangential distortion coefficient.
+            s1: First thin-prism distortion coefficient.
+            s2: Second thin-prism distortion coefficient.
+            s3: Third thin-prism distortion coefficient.
+            s4: Fourth thin-prism distortion coefficient.
+            out_rays: Optional output array to write into, shape
+                ``(out_camera_count, height, width, 2)``. If ``None``,
+                allocates a new array.
+            camera_index: Camera index in *out_rays* to write. Ignored when
+                *out_rays* is ``None``.
+
+        Returns:
+            camera_rays: *out_rays* if provided, otherwise a new array with
+                shape ``(1, height, width, 2)`` and dtype ``vec3f``.
+
+        Raises:
+            ValueError: If any calibration value is non-finite, or if focal
+                lengths or calibration image dimensions are non-positive.
+        """
+        image_width = width if image_width is None else image_width
+        image_height = height if image_height is None else image_height
+        if not (math.isfinite(fx) and math.isfinite(fy) and fx > 0.0 and fy > 0.0):
+            raise ValueError("fx and fy must be finite and positive.")
+        if not (
+            math.isfinite(image_width) and math.isfinite(image_height) and image_width > 0.0 and image_height > 0.0
+        ):
+            raise ValueError("image_width and image_height must be finite and positive.")
+        if not all(math.isfinite(value) for value in (cx, cy, k1, k2, k3, k4, k5, k6, p1, p2, s1, s2, s3, s4)):
+            raise ValueError("cx, cy, and distortion coefficients must be finite.")
+        out_rays, camera_index = camera_utils._validate_camera_ray_output(
+            width, height, 1, out_rays, camera_index, self.__render_context.device
+        )
+
+        wp.launch(
+            kernel=camera_utils.compute_camera_rays_pinhole_opencv_kernel,
+            dim=(height, width),
+            inputs=[
+                width,
+                height,
+                image_width,
+                image_height,
+                fx,
+                fy,
+                cx,
+                cy,
+                k1,
+                k2,
+                k3,
+                k4,
+                k5,
+                k6,
+                p1,
+                p2,
+                s1,
+                s2,
+                s3,
+                s4,
+                camera_index,
+                out_rays,
+            ],
+            device=self.__render_context.device,
+        )
+
+        return out_rays
+
     def compute_camera_rays_fisheye_opencv(
         self,
         width: int,