gsplat 1.0.0

data/lib/gsplat/backend/ruby/projection.rb

@@ -0,0 +1,251 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Numo implementation of dense fused Gaussian projection.
+    module RubyProjection
+      module_function
+
+      # rubocop:disable Metrics/ParameterLists
+      def forward(means, covars, quaternions, scales, viewmats, intrinsics, width, height,
+                  eps2d:, near_plane:, far_plane:, radius_clip:, calc_compensations:, camera_model:,
+                  radial_coeffs: nil, tangential_coeffs: nil, thin_prism_coeffs: nil,
+                  global_z_order: true)
+        # rubocop:enable Metrics/ParameterLists
+        inputs = prepare_inputs(
+          means, covars, quaternions, scales, viewmats, intrinsics, width, height,
+          eps2d, near_plane, far_plane, radius_clip, camera_model,
+          radial_coeffs, tangential_coeffs, thin_prism_coeffs
+        )
+        means, covars, viewmats, intrinsics = inputs
+        camera_means, camera_covars = Math::CameraProjection.world_to_cam(means, covars, viewmats)
+        projection_means = projection_safe_means(camera_means, near_plane, far_plane)
+        means2d, covars2d = project(
+          projection_means,
+          camera_covars,
+          intrinsics,
+          width,
+          height,
+          camera_model,
+          radial_coeffs: radial_coeffs,
+          tangential_coeffs: tangential_coeffs,
+          thin_prism_coeffs: thin_prism_coeffs
+        )
+        finish_projection(
+          camera_means,
+          means2d,
+          covars2d,
+          width,
+          height,
+          eps2d,
+          near_plane,
+          far_plane,
+          radius_clip,
+          calc_compensations,
+          global_z_order: global_z_order
+        )
+      end
+
+      # rubocop:disable Metrics/AbcSize, Metrics/ParameterLists
+      def prepare_inputs(means, covars, quaternions, scales, viewmats, intrinsics, width, height,
+                         eps2d, near_plane, far_plane, radius_clip, camera_model,
+                         radial_coeffs = nil, tangential_coeffs = nil, thin_prism_coeffs = nil)
+        # rubocop:enable Metrics/AbcSize, Metrics/ParameterLists
+        unless means.is_a?(Numo::NArray) && [Numo::SFloat, Numo::DFloat].include?(means.class)
+          raise ArgumentError, "means must be Numo::SFloat or Numo::DFloat"
+        end
+
+        valid_means = means.ndim == 2 && means.shape[-1] == 3
+        raise ShapeError, "expected means [N,3], got #{means.shape.inspect}" unless valid_means
+
+        covars = prepare_covariances(means, covars, quaternions, scales)
+        viewmats = means.class.cast(viewmats)
+        intrinsics = means.class.cast(intrinsics)
+        unless viewmats.ndim == 3 && viewmats.shape[1..] == [4, 4]
+          raise ShapeError, "expected viewmats [C,4,4], got #{viewmats.shape.inspect}"
+        end
+        unless intrinsics.shape == [viewmats.shape[0], 3, 3]
+          raise ShapeError, "expected intrinsics [#{viewmats.shape[0]},3,3], got #{intrinsics.shape.inspect}"
+        end
+
+        validate_options!(
+          width: width,
+          height: height,
+          eps2d: eps2d,
+          near_plane: near_plane,
+          far_plane: far_plane,
+          radius_clip: radius_clip,
+          camera_model: camera_model
+        )
+        validate_distortion!(
+          viewmats.shape[0], camera_model,
+          radial_coeffs, tangential_coeffs, thin_prism_coeffs
+        )
+        [means, covars, viewmats, intrinsics]
+      end
+
+      def prepare_covariances(means, covars, quaternions, scales)
+        if covars
+          raise ArgumentError, "provide covars or quaternions/scales, not both" if quaternions || scales
+
+          return expand_covariances(means.class.cast(covars), means.shape[0])
+        end
+        raise ArgumentError, "quaternions and scales are required when covars is nil" unless quaternions && scales
+
+        covariance, = RubyQuatScaleToCovarPreci.forward(
+          means.class.cast(quaternions),
+          means.class.cast(scales),
+          compute_covar: true,
+          compute_preci: false,
+          triu: false
+        )
+        covariance
+      end
+      private_class_method :prepare_covariances
+
+      def expand_covariances(covars, gaussian_count)
+        return covars if covars.shape == [gaussian_count, 3, 3]
+        unless covars.shape == [gaussian_count, 6]
+          raise ShapeError, "expected covars [N,3,3] or [N,6], got #{covars.shape.inspect}"
+        end
+
+        output = covars.class.zeros(gaussian_count, 3, 3)
+        [[0, 0, 0], [0, 1, 1], [0, 2, 2], [1, 1, 3], [1, 2, 4], [2, 2, 5]].each do |row, column, index|
+          output[true, row, column] = covars[true, index]
+          output[true, column, row] = covars[true, index]
+        end
+        output
+      end
+      private_class_method :expand_covariances
+
+      # rubocop:disable Metrics/ParameterLists
+      def validate_options!(width:, height:, eps2d:, near_plane:, far_plane:, radius_clip:, camera_model:)
+        # rubocop:enable Metrics/ParameterLists
+        unless width.is_a?(Integer) && width.positive? && height.is_a?(Integer) && height.positive?
+          raise ArgumentError, "width and height must be positive integers"
+        end
+        raise ArgumentError, "eps2d and radius_clip must be non-negative" if eps2d.negative? || radius_clip.negative?
+        raise ArgumentError, "near_plane must be less than far_plane" unless near_plane < far_plane
+        return if %w[pinhole ortho fisheye].include?(camera_model.to_s)
+
+        raise ArgumentError, "camera_model must be \"pinhole\", \"ortho\", or \"fisheye\""
+      end
+      private_class_method :validate_options!
+
+      def validate_distortion!(camera_count, camera_model, radial, tangential, thin_prism)
+        if camera_model.to_s == "ortho" && [radial, tangential, thin_prism].any?
+          raise ArgumentError, "distortion coefficients are unsupported for orthographic cameras"
+        end
+
+        expected_radial = camera_model.to_s == "fisheye" ? [4] : [4, 6]
+        validate_coefficients!(radial, camera_count, expected_radial, "radial_coeffs")
+        validate_coefficients!(tangential, camera_count, [2], "tangential_coeffs")
+        validate_coefficients!(thin_prism, camera_count, [4], "thin_prism_coeffs")
+      end
+      private_class_method :validate_distortion!
+
+      def validate_coefficients!(value, camera_count, widths, name)
+        return unless value
+        unless value.is_a?(Numo::NArray) && [Numo::SFloat, Numo::DFloat].include?(value.class) &&
+               value.ndim == 2 && value.shape[0] == camera_count && widths.include?(value.shape[1])
+          raise ShapeError, "expected #{name} [#{camera_count},#{widths.join(' or ')}]"
+        end
+      end
+      private_class_method :validate_coefficients!
+
+      # rubocop:disable Metrics/ParameterLists
+      def project(means, covars, intrinsics, width, height, camera_model,
+                  radial_coeffs: nil, tangential_coeffs: nil, thin_prism_coeffs: nil)
+        # rubocop:enable Metrics/ParameterLists
+        if camera_model.to_s == "ortho"
+          return Math::CameraProjection.ortho_proj(means, covars, intrinsics, width, height)
+        end
+        if camera_model.to_s == "fisheye" || [radial_coeffs, tangential_coeffs, thin_prism_coeffs].any?
+          return Math::CameraProjection.distorted_proj(
+            means,
+            covars,
+            intrinsics,
+            camera_model,
+            radial_coeffs: radial_coeffs,
+            tangential_coeffs: tangential_coeffs,
+            thin_prism_coeffs: thin_prism_coeffs
+          )
+        end
+
+        Math::CameraProjection.persp_proj(means, covars, intrinsics, width, height)
+      end
+
+      def projection_safe_means(camera_means, near_plane, far_plane)
+        output = camera_means.dup
+        depths = output[true, true, 2]
+        culled = depths.le(near_plane) | depths.ge(far_plane)
+        depths[culled] = 1 if culled.any?
+        output[true, true, 2] = depths
+        output
+      end
+      private_class_method :projection_safe_means
+
+      # rubocop:disable Metrics/AbcSize, Metrics/ParameterLists
+      def finish_projection(camera_means, means2d, covars2d, width, height, eps2d,
+                            near_plane, far_plane, radius_clip, calc_compensations,
+                            global_z_order: true)
+        # rubocop:enable Metrics/ParameterLists
+        determinant_original = Math::Mat.det2x2(covars2d)
+        regularized = covars2d.dup
+        regularized[true, true, 0, 0] += eps2d
+        regularized[true, true, 1, 1] += eps2d
+        determinant = Math::Mat.det2x2(regularized)
+        safe_determinant = determinant.dup
+        invalid_determinant = safe_determinant.le(0)
+        safe_determinant[invalid_determinant] = 1e-10 if invalid_determinant.any?
+        conics = conics_from_covariance(regularized, safe_determinant)
+        compensations = compensation(determinant_original, safe_determinant) if calc_compensations
+        variance_x = regularized[true, true, 0, 0].dup
+        variance_y = regularized[true, true, 1, 1].dup
+        variance_x[variance_x.lt(0)] = 0 if variance_x.lt(0).any?
+        variance_y[variance_y.lt(0)] = 0 if variance_y.lt(0).any?
+        radius_x = Numo::Int32.cast((3.33 * (variance_x**0.5)).ceil)
+        radius_y = Numo::Int32.cast((3.33 * (variance_y**0.5)).ceil)
+        z_depths = camera_means[true, true, 2]
+        depths = if global_z_order
+                   z_depths.dup
+                 else
+                   (camera_means**2).sum(axis: 2)**0.5
+                 end
+        visible = determinant.gt(0) & z_depths.gt(near_plane) & z_depths.lt(far_plane)
+        visible &= radius_x.gt(radius_clip) | radius_y.gt(radius_clip)
+        visible &= (means2d[true, true, 0] + radius_x).gt(0)
+        visible &= means2d[true, true, 0] - radius_x < width
+        visible &= (means2d[true, true, 1] + radius_y).gt(0)
+        visible &= means2d[true, true, 1] - radius_y < height
+        radius_x[visible.eq(0)] = 0
+        radius_y[visible.eq(0)] = 0
+        radii = Numo::Int32.zeros(*(depths.shape + [2]))
+        radii[true, true, 0] = radius_x
+        radii[true, true, 1] = radius_y
+        [radii, means2d, depths, conics, compensations]
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :finish_projection
+
+      def conics_from_covariance(covariance, determinant)
+        output = covariance.class.zeros(*(covariance.shape[0...-2] + [3]))
+        output[true, true, 0] = covariance[true, true, 1, 1] / determinant
+        output[true, true, 1] = -(
+          covariance[true, true, 0, 1] + covariance[true, true, 1, 0]
+        ) / (2.0 * determinant)
+        output[true, true, 2] = covariance[true, true, 0, 0] / determinant
+        output
+      end
+      private_class_method :conics_from_covariance
+
+      def compensation(original, regularized)
+        ratio = original / regularized
+        negative = ratio.lt(0)
+        ratio[negative] = 0 if negative.any?
+        ratio**0.5
+      end
+      private_class_method :compensation
+    end
+  end
+end

data/lib/gsplat/backend/ruby/projection_backward.rb

@@ -0,0 +1,190 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Analytic Numo VJP for dense pinhole projection.
+    module RubyProjectionBackward
+      module_function
+
+      # rubocop:disable Metrics/ParameterLists
+      def backward(means, covars, quaternions, scales, viewmats, intrinsics, width, height,
+                   grad_means2d, grad_depths, grad_conics, grad_compensations, **options)
+        # rubocop:enable Metrics/ParameterLists
+        prepared = RubyProjection.prepare_inputs(
+          means, covars, quaternions, scales, viewmats, intrinsics, width, height,
+          options.fetch(:eps2d), options.fetch(:near_plane), options.fetch(:far_plane),
+          options.fetch(:radius_clip), options.fetch(:camera_model),
+          options[:radial_coeffs], options[:tangential_coeffs], options[:thin_prism_coeffs]
+        )
+        means, expanded_covars, viewmats, intrinsics = prepared
+        camera_means, camera_covars = Math::CameraProjection.world_to_cam(means, expanded_covars, viewmats)
+        projection_means = safe_means(
+          camera_means,
+          options.fetch(:near_plane),
+          options.fetch(:far_plane)
+        )
+        camera_model = options.fetch(:camera_model)
+        means2d, covars2d = RubyProjection.project(
+          projection_means,
+          camera_covars,
+          intrinsics,
+          width,
+          height,
+          camera_model,
+          radial_coeffs: options[:radial_coeffs],
+          tangential_coeffs: options[:tangential_coeffs],
+          thin_prism_coeffs: options[:thin_prism_coeffs]
+        )
+        jacobians = projection_jacobians(projection_means, intrinsics, width, height, options)
+        output_gradients = validate_output_gradients(
+          means.class,
+          means2d,
+          grad_means2d,
+          grad_depths,
+          grad_conics,
+          grad_compensations
+        )
+        grad_camera_means, grad_camera_covars = projection_vjp(
+          projection_means,
+          camera_covars,
+          covars2d,
+          jacobians,
+          intrinsics,
+          output_gradients,
+          options.fetch(:eps2d),
+          width,
+          height,
+          options
+        )
+        grad_means, grad_expanded_covars = world_to_cam_vjp(
+          grad_camera_means,
+          grad_camera_covars,
+          viewmats
+        )
+        input_covariance_gradients(
+          means,
+          covars,
+          quaternions,
+          scales,
+          grad_means,
+          grad_expanded_covars
+        )
+      end
+
+      def safe_means(camera_means, near_plane, far_plane)
+        output = camera_means.dup
+        depths = output[true, true, 2]
+        culled = depths.le(near_plane) | depths.ge(far_plane)
+        depths[culled] = 1 if culled.any?
+        output[true, true, 2] = depths
+        output
+      end
+      private_class_method :safe_means
+
+      def projection_jacobians(means, intrinsics, width, height, options)
+        output = means.class.zeros(means.shape[0], means.shape[1], 2, 3)
+        means.shape[0].times do |camera_index|
+          camera_means = means[camera_index, true, true]
+          camera_intrinsics = intrinsics[camera_index, true, true]
+          camera_model = options.fetch(:camera_model)
+          extended = camera_model.to_s == "fisheye" ||
+                     %i[radial_coeffs tangential_coeffs thin_prism_coeffs].any? { |key| options[key] }
+          _, jacobians = if camera_model.to_s == "ortho"
+                           Math::CameraProjection.ortho_camera(camera_means, camera_intrinsics)
+                         elsif extended
+                           Math::CameraDistortion.project_camera(
+                             camera_means,
+                             camera_intrinsics,
+                             camera_model,
+                             radial: coefficient_row(options[:radial_coeffs], camera_index),
+                             tangential: coefficient_row(options[:tangential_coeffs], camera_index),
+                             thin_prism: coefficient_row(options[:thin_prism_coeffs], camera_index)
+                           )
+                         else
+                           Math::CameraProjection.pinhole_camera(
+                             camera_means, camera_intrinsics, width, height
+                           )
+                         end
+          output[camera_index, true, true, true] = jacobians
+        end
+        output
+      end
+      private_class_method :projection_jacobians
+
+      def coefficient_row(value, camera_index)
+        value && value[camera_index, true].to_a
+      end
+      private_class_method :coefficient_row
+
+      # rubocop:disable Metrics/ParameterLists
+      def validate_output_gradients(type, means2d, grad_means2d, grad_depths, grad_conics, grad_compensations)
+        # rubocop:enable Metrics/ParameterLists
+        leading_shape = means2d.shape[0...-1]
+        [
+          cast_gradient(type, grad_means2d, means2d.shape),
+          cast_gradient(type, grad_depths, leading_shape),
+          cast_gradient(type, grad_conics, leading_shape + [3]),
+          grad_compensations && cast_gradient(type, grad_compensations, leading_shape)
+        ]
+      end
+      private_class_method :validate_output_gradients
+
+      def cast_gradient(type, gradient, expected_shape)
+        value = type.cast(gradient)
+        unless value.shape == expected_shape
+          raise ShapeError, "expected output gradient #{expected_shape.inspect}, got #{value.shape.inspect}"
+        end
+
+        value
+      end
+      private_class_method :cast_gradient
+
+      # rubocop:disable Metrics/ParameterLists
+      def projection_vjp(means, covars, covars2d, jacobians, intrinsics, gradients, eps2d, width, height,
+                         options)
+        # rubocop:enable Metrics/ParameterLists
+        grad_means2d, grad_depths, grad_conics, grad_compensations = gradients
+        grad_covars2d = covariance_output_vjp(
+          covars2d,
+          grad_conics,
+          grad_compensations,
+          eps2d
+        )
+        count = means.shape[0] * means.shape[1]
+        jacobian_flat = jacobians.reshape(count, 2, 3)
+        covariance_flat = covars.reshape(count, 3, 3)
+        gradient_flat = grad_covars2d.reshape(count, 2, 2)
+        grad_camera_covars = Math::Mat.matmul_batch(
+          Math::Mat.matmul_batch(jacobian_flat.transpose(0, 2, 1), gradient_flat),
+          jacobian_flat
+        ).reshape(*covars.shape)
+        grad_jacobians = jacobian_vjp(jacobian_flat, covariance_flat, gradient_flat)
+        grad_camera_means = camera_mean_vjp(
+          means,
+          intrinsics,
+          grad_means2d,
+          grad_depths,
+          grad_jacobians.reshape(*jacobians.shape),
+          width,
+          height,
+          options
+        )
+        [grad_camera_means, grad_camera_covars]
+      end
+      private_class_method :projection_vjp
+
+      def jacobian_vjp(jacobians, covars, gradient)
+        first = Math::Mat.matmul_batch(
+          Math::Mat.matmul_batch(gradient, jacobians),
+          covars.transpose(0, 2, 1)
+        )
+        second = Math::Mat.matmul_batch(
+          Math::Mat.matmul_batch(gradient.transpose(0, 2, 1), jacobians),
+          covars
+        )
+        first + second
+      end
+      private_class_method :jacobian_vjp
+    end
+  end
+end

data/lib/gsplat/backend/ruby/projection_covariance_vjp.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Covariance-output VJP helpers for RubyProjectionBackward.
+    module RubyProjectionBackward
+      module_function
+
+      # rubocop:disable Metrics/AbcSize
+      def covariance_output_vjp(covars2d, grad_conics, grad_compensations, eps2d)
+        regularized = covars2d.dup
+        regularized[true, true, 0, 0] += eps2d
+        regularized[true, true, 1, 1] += eps2d
+        grad_regularized = conic_vjp(regularized, grad_conics)
+        return grad_regularized unless grad_compensations
+
+        original_det = Math::Mat.det2x2(covars2d)
+        regularized_det = Math::Mat.det2x2(regularized)
+        ratio = original_det / regularized_det
+        positive = ratio.gt(0)
+        compensation = ratio.class.zeros(*ratio.shape)
+        compensation[positive] = ratio[positive]**0.5 if positive.any?
+        grad_original_det = ratio.class.zeros(*ratio.shape)
+        if positive.any?
+          grad_original_det[positive] = grad_compensations[positive] /
+                                        (2 * compensation[positive] * regularized_det[positive])
+        end
+        grad_regularized_det = -grad_compensations * compensation / (2 * regularized_det)
+        grad_regularized + determinant_vjp(regularized, grad_regularized_det) +
+          determinant_vjp(covars2d, grad_original_det)
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :covariance_output_vjp
+
+      # rubocop:disable Metrics/AbcSize
+      def conic_vjp(covariance, gradient)
+        count = gradient.size / 3
+        matrices = covariance.reshape(count, 2, 2)
+        determinant = Math::Mat.det2x2(matrices)
+        inverse = matrices.class.zeros(count, 2, 2)
+        inverse[true, 0, 0] = matrices[true, 1, 1] / determinant
+        inverse[true, 0, 1] = -matrices[true, 0, 1] / determinant
+        inverse[true, 1, 0] = -matrices[true, 1, 0] / determinant
+        inverse[true, 1, 1] = matrices[true, 0, 0] / determinant
+        grad_flat = gradient.reshape(count, 3)
+        grad_inverse = matrices.class.zeros(count, 2, 2)
+        grad_inverse[true, 0, 0] = grad_flat[true, 0]
+        grad_inverse[true, 0, 1] = grad_flat[true, 1] * 0.5
+        grad_inverse[true, 1, 0] = grad_flat[true, 1] * 0.5
+        grad_inverse[true, 1, 1] = grad_flat[true, 2]
+        inverse_transpose = inverse.transpose(0, 2, 1)
+        result = -Math::Mat.matmul_batch(
+          Math::Mat.matmul_batch(inverse_transpose, grad_inverse),
+          inverse_transpose
+        )
+        result.reshape(*covariance.shape)
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :conic_vjp
+
+      def determinant_vjp(matrix, gradient)
+        output = matrix.class.zeros(*matrix.shape)
+        output[true, true, 0, 0] = gradient * matrix[true, true, 1, 1]
+        output[true, true, 0, 1] = -gradient * matrix[true, true, 1, 0]
+        output[true, true, 1, 0] = -gradient * matrix[true, true, 0, 1]
+        output[true, true, 1, 1] = gradient * matrix[true, true, 0, 0]
+        output
+      end
+      private_class_method :determinant_vjp
+    end
+  end
+end