3D Gaussian Splatting (3DGS) Implementation in PyTorch

This implementation builds 3D Gaussian Splatting from scratch. It includes per-Gaussian learnable positions, scales, quaternions and RGBA color, a 3D covariance computed as R S Sᵀ Rᵀ, projection to 2D via a focal-length perspective Jacobian, and splat-based alpha compositing over a pixel grid. A smoke test verifies initialization shapes, covariance symmetry, quaternion-to-rotation orthogonality, and that the rendered image stays inside [0, 1].

import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Optional

class GaussianSplatting3D:
    """
    - Position (x, y, z)
    - Scale (sx, sy, sz)
    - Rotation (quaternion)
    - Color (RGB + opacity)
    """
    def __init__(self, num_points: int, device: str = 'cuda' if torch.cuda.is_available() else 'cpu'):

        self.device = device
        self.num_points = num_points
        
        self.positions = nn.Parameter(torch.randn(num_points, 3, device=device))
        self.scales = nn.Parameter(torch.ones(num_points, 3, device=device))
        self.rotations = nn.Parameter(torch.zeros(num_points, 4, device=device))
        self.rotations.data[:, 0] = 1.0  # identity rotation
        self.colors = nn.Parameter(torch.ones(num_points, 4, device=device))

    def get_covariance_matrices(self) -> torch.Tensor:
        quats = self.rotations / torch.norm(self.rotations, dim=1, keepdim=True)
        
        # Convert quaternions to rotation matrices
        R = self._quaternion_to_rotation_matrix(quats)
        
        # Create scale matrices 
        S = torch.diag_embed(self.scales ** 2)
        
        # Covariance = R * S * S * R^T
        covariance = torch.matmul(torch.matmul(R, S), R.transpose(1, 2))
        return covariance

    def _quaternion_to_rotation_matrix(self, q: torch.Tensor) -> torch.Tensor:
        w, x, y, z = q.unbind(1)
        
        rot = torch.stack([
            1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y,
            2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x,
            2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y
        ], dim=1).reshape(-1, 3, 3)
        
        return rot

    def render(self, camera_pos: torch.Tensor, camera_rot: torch.Tensor, 
               image_size: Tuple[int, int], focal_length: float) -> torch.Tensor:
        """  
        Args:
            camera_pos: Camera position in world space
            camera_rot: Camera rotation matrix
            image_size: Output image dimensions 
            focal_length: Camera focal length
        """
        height, width = image_size

        points_camera = torch.matmul(camera_rot, self.positions.T).T + camera_pos
        points_image = self._project_points(points_camera, focal_length)
        covs_3d = self.get_covariance_matrices()
        covs_2d = self._project_covariance(covs_3d, points_camera, camera_rot, focal_length)

        image = torch.zeros(height, width, 3, device=self.device)

        for i in range(self.num_points):
            if points_camera[i, 2] <= 0:
                continue

            bounds = self._get_gaussian_bounds(points_image[i], covs_2d[i], image_size)
            if bounds is None:
                continue

            min_x, max_x, min_y, max_y = bounds

            # Updated meshgrid 
            y_coords, x_coords = torch.meshgrid(
                torch.arange(min_y, max_y, device=self.device),
                torch.arange(min_x, max_x, device=self.device),
                indexing='ij' 
            )
            coords = torch.stack([x_coords, y_coords], dim=-1)

            gaussian = self._evaluate_gaussian_2d(coords, points_image[i], covs_2d[i])
            
            gaussian = gaussian.unsqueeze(-1)  
            color = self.colors[i, :3].view(1, 1, 3)  
            alpha = self.colors[i, 3]

            # Calculate color contribution
            color_contribution = gaussian * color

            # Update image region
            current_region = image[min_y:max_y, min_x:max_x]
            image[min_y:max_y, min_x:max_x] = current_region * (1 -  gaussian * alpha) + color_contribution * alpha

        return image

    def _project_points(self, points_camera: torch.Tensor, focal_length: float) -> torch.Tensor:
        """Project 3D points to 2D image plane."""
        points_image = torch.zeros_like(points_camera[:, :2])
        points_image[:, 0] = focal_length * points_camera[:, 0] / points_camera[:, 2]
        points_image[:, 1] = focal_length * points_camera[:, 1] / points_camera[:, 2]
        return points_image

    def _project_covariance(self, covs_3d: torch.Tensor, points_camera: torch.Tensor,
                          camera_rot: torch.Tensor, focal_length: float) -> torch.Tensor:
        """Project 3D covariance matrices to 2D image plane."""
        covs_camera = torch.matmul(
            torch.matmul(camera_rot, covs_3d),
            camera_rot.transpose(0, 1)
        )

        # Create projection jacobian
        z = points_camera[:, 2:3]
        J = torch.zeros(self.num_points, 2, 3, device=self.device)
        J[:, 0, 0] = focal_length / z.squeeze()
        J[:, 1, 1] = focal_length / z.squeeze()
        J[:, 0, 2] = -focal_length * points_camera[:, 0] / (z * z).squeeze()
        J[:, 1, 2] = -focal_length * points_camera[:, 1] / (z * z).squeeze()

        # Project to 2D using J * Σ * J^T
        covs_2d = torch.matmul(torch.matmul(J, covs_camera), J.transpose(1, 2))
        return covs_2d

    def _get_gaussian_bounds(self, point_2d: torch.Tensor, cov_2d: torch.Tensor,
                           image_size: Tuple[int, int]) -> Optional[Tuple[int, int, int, int]]:
        """Calculate pixel bounds."""
        height, width = image_size
        
        std = torch.sqrt(torch.diagonal(cov_2d))
        radius = 3 * torch.max(std)
        
        # Calculate bounds
        min_x = max(0, int(point_2d[0] - radius))
        max_x = min(width, int(point_2d[0] + radius + 1))
        min_y = max(0, int(point_2d[1] - radius))
        max_y = min(height, int(point_2d[1] + radius + 1))
        
        if min_x >= max_x or min_y >= max_y:
            return None
            
        return min_x, max_x, min_y, max_y

    def _evaluate_gaussian_2d(self, coords: torch.Tensor, mean: torch.Tensor,
                            cov: torch.Tensor) -> torch.Tensor:
        """Evaluate 2D Gaussian at given coordinates."""
        diff = coords - mean.unsqueeze(0).unsqueeze(0)
        
        # Compute mahalanobis distance
        inv_cov = torch.inverse(cov)
        mahalanobis = torch.sum(
            torch.matmul(diff, inv_cov) * diff,
            dim=-1
        )
        
        gaussian = torch.exp(-0.5 * mahalanobis)
        gaussian = gaussian / (2 * np.pi * torch.sqrt(torch.det(cov)))
        
        return gaussian

def run_tests():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Running tests on device: {device}")

    print("\nTesting initialization...")
    num_points = 100
    gs = GaussianSplatting3D(num_points, device)
    assert gs.positions.shape == (num_points, 3), f"Expected positions shape (100, 3), got {gs.positions.shape}"
    assert gs.scales.shape == (num_points, 3), f"Expected scales shape (100, 3), got {gs.scales.shape}"
    assert gs.rotations.shape == (num_points, 4), f"Expected rotations shape (100, 4), got {gs.rotations.shape}"
    assert gs.colors.shape == (num_points, 4), f"Expected colors shape (100, 4), got {gs.colors.shape}"
    print("✓ Initialization test passed")

    print("\nTesting covariance computation...")
    covs = gs.get_covariance_matrices()
    assert covs.shape == (num_points, 3, 3), f"Expected covariance shape (100, 3, 3), got {covs.shape}"
    symmetry_diff = torch.max(torch.abs(covs - covs.transpose(1, 2)))
    assert symmetry_diff < 1e-5, f"Covariance matrices not symmetric, max difference: {symmetry_diff}"
    print("✓ Covariance computation test passed")

    print("\nTesting quaternion conversion...")
    identity_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=device)
    rot_mat = gs._quaternion_to_rotation_matrix(identity_quat)
    identity_error = torch.max(torch.abs(rot_mat[0] - torch.eye(3, device=device)))
    assert identity_error < 1e-5, f"Identity quaternion did not produce identity matrix, max error: {identity_error}"
    
    orthogonality_error = torch.max(torch.abs(
        torch.matmul(rot_mat, rot_mat.transpose(1, 2)) - torch.eye(3, device=device).unsqueeze(0)
    ))
    assert orthogonality_error < 1e-5, f"Rotation matrix not orthogonal, max error: {orthogonality_error}"
    print("✓ Quaternion conversion test passed")

    print("\nTesting rendering...")
    camera_pos = torch.zeros(3, device=device)
    camera_rot = torch.eye(3, device=device)
    image_size = (480, 640)
    image = gs.render(
        camera_pos,
        camera_rot,
        image_size=image_size,
        focal_length=500.0
    )
    
    assert image.shape == (480, 640, 3), f"Expected image shape (480, 640, 3), got {image.shape}"
    assert torch.all(image >= 0), "Image contains negative values"
    assert torch.all(image <= 1), "Image contains values greater than 1"
    print("✓ Rendering test passed")

    print("\nAll tests passed successfully!")

if __name__ == "__main__":
    run_tests()