Depth returned by the gs-madrona batch renderer with 'use_rasterizer=True' differed from depth returned by the rasterizer

[Ella-e]

I used the branch support-batch-renderer together with gs-madrona to render rgba and depth images in batches. I observed that with the same camera settings, the depth returned by

gs.options.renderers.BatchRenderer(
    use_rasterizer=True,
)

is different from the depth returned by the rasterizer renderer (self._rasterizer.render_camera()) from the main branch.
As a result, if I use a similar way to construct the point cloud from depth as in camera.render_pointcloud(), there will be an offset from the world frame. The further the camera is placed away from the center, the larger the offset is. Therefore, I am wondering if the depth returned by the MadronaBatchRenderer in rasterizer mode is calculated in a different way than that of the rasterizer renderer in the main branch, and should be treated differently when constructing point cloud?

The offset looks like this:
This image contains 5 point clouds rendered by 5 different cameras. They are supposed to be shown as a single point cloud after fusion. My method for construct the piont cloud and fusion works fine for the depth returned by self._rasterizer.render_camera().

[duburcqa]

@yuhongyi @ACMLCZH

[yuhongyi]

This is expected, since the depth buffer in the gs-madrona rasterizer pipeline is non-linear. Some postprocessing is done to reconstruct the linear depth, but it is not as accurate. There is a way to fix this. We will evaluate the solution.

[Ella-e]

I see, thank you very much!

[ACMLCZH]

Does the depth image rendered by use_rasterizer=False align with depth return by the rasterizer?

[Ella-e]

With use_rasterizer=False, the depth values returned is likely the Euclidean distance between the point and camera's optical center, rather than the linear depth returned by the rasterizer. However, with use_rasterizer=False, the calculated point cloud doesn't have offset and looks good if the returned depth values are handled correctly.

Here is the code I used for calculating the points for the point cloud.

v_coords, u_coords = torch.meshgrid(
    torch.arange(H, device=device, dtype=torch.float32),
    torch.arange(W, device=device, dtype=torch.float32),
    indexing='ij'
) # (H, W)
u_coords = u_coords.view(1, 1, H, W)
v_coords = v_coords.view(1, 1, H, W) 
fx = K_batch[:, :, 0, 0].view(n_env, n_cam, 1, 1)
fy = K_batch[:, :, 1, 1].view(n_env, n_cam, 1, 1)
cx = K_batch[:, :, 0, 2].view(n_env, n_cam, 1, 1)
cy = K_batch[:, :, 1, 2].view(n_env, n_cam, 1, 1)

# With `use_rasterizer=True`
# z_cam = depth # (n_env, n_cam, H, W)
# x_cam = (u_coords - cx) * z_cam / fx
# y_cam = (v_coords - cy) * z_cam / fy`

# With `use_rasterizer=False`
x_nd = (u_coords - cx) / fx
y_nd = (v_coords - cy) / fy
den = torch.sqrt(x_nd**2 + y_nd**2 + 1.0)
z_cam = depth / den
x_cam = x_nd * z_cam
y_cam = y_nd * z_cam

pc_cam = torch.stack([x_cam, y_cam, z_cam], dim=-1)

[duburcqa]

@Ella-e Could you share the entire code? I would love to make a unit test out of it!

[Ella-e]

Sure, here is the complete callable function that construct point cloud in batch.

import torch
import numpy as np
import math
from scipy.spatial.transform import Rotation

def get_pointcloud_batch(
    depth_tuple: tuple,
    camera_pos: torch.Tensor,
    camera_quat: torch.Tensor,
    camera_fov: tuple,
    znear: float,
    zfar: float,
    world_frame: bool = True
):
    device = camera_pos.device
    
    # Stack the depth tuple into a single tensor and remove the last dimension
    depth = torch.stack([d.squeeze(-1) for d in depth_tuple], dim=1).to(device)

    n_env, n_cam, H, W = depth.shape
    
    # Camera Intrinsics
    K_batch = torch.zeros(n_env, n_cam, 3, 3, device=device)
    fov_rad = torch.tensor([math.radians(fov) for fov in camera_fov], device=device) # (n_cam,)
    fy = (H / 2) / torch.tan(fov_rad / 2) # (n_cam,)
    fx = fy.clone() # Assuming square pixels
    
    K_cam = torch.zeros(n_cam, 3, 3, device=device)
    K_cam[:, 0, 0] = fx
    K_cam[:, 1, 1] = fy
    K_cam[:, 0, 2] = W / 2
    K_cam[:, 1, 2] = H / 2
    K_cam[:, 2, 2] = 1
    
    K_batch = K_cam.unsqueeze(0).expand(n_env, -1, -1, -1) # (n_env, n_cam, 3, 3)
    
    # Camera Extrinsics (Convert from wxyz to xyzw first)
    quat_xyzw = camera_quat[..., [1, 2, 3, 0]].cpu().numpy()
    rot_matrix_np = Rotation.from_quat(quat_xyzw.reshape(-1, 4)).as_matrix()
    fix_rotation = Rotation.from_euler('x', 90, degrees=True) # From z-up to y-up
    fixed_rot_matrix_np = rot_matrix_np @ fix_rotation.as_matrix()
    rot_matrix = torch.from_numpy(fixed_rot_matrix_np).to(device) # (n_env * n_cam, 3, 3)
    rot_matrix = rot_matrix.reshape(n_env, n_cam, 3, 3) # (n_env, n_cam, 3, 3)
    
    # OpenGL to World
    pose_opengl_to_world = torch.eye(4, device=device).unsqueeze(0).unsqueeze(0).repeat(n_env, n_cam, 1, 1)
    pose_opengl_to_world[:, :, :3, :3] = rot_matrix
    pose_opengl_to_world[:, :, :3, 3] = camera_pos
    
    T_OPENGL_TO_OPENCV = torch.tensor([
        [1, 0, 0, 0],
        [0, -1, 0, 0],
        [0, 0, -1, 0],
        [0, 0, 0, 1]
    ], dtype=torch.float32, device=device)
    
    pose_opencv_to_world = pose_opengl_to_world @ T_OPENGL_TO_OPENCV
    
    # Pixel Grid Mesh
    v_coords, u_coords = torch.meshgrid(
        torch.arange(H, device=device, dtype=torch.float32),
        torch.arange(W, device=device, dtype=torch.float32),
        indexing='ij'
    ) # (H, W)
    
    # (1, 1, H, W)
    u_coords = u_coords.view(1, 1, H, W)
    v_coords = v_coords.view(1, 1, H, W)
    
    fx = K_batch[:, :, 0, 0].view(n_env, n_cam, 1, 1)
    fy = K_batch[:, :, 1, 1].view(n_env, n_cam, 1, 1)
    cx = K_batch[:, :, 0, 2].view(n_env, n_cam, 1, 1)
    cy = K_batch[:, :, 1, 2].view(n_env, n_cam, 1, 1)
    
    x_nd = (u_coords - cx) / fx
    y_nd = (v_coords - cy) / fy
    den = torch.sqrt(x_nd**2 + y_nd**2 + 1.0)
    z_cam = depth / den
    
    x_cam = x_nd * z_cam
    y_cam = y_nd * z_cam
    
    # Use Rasterizer
    # z_cam = depth # (n_env, n_cam, H, W)
    # x_cam = (u_coords - cx) * z_cam / fx
    # y_cam = (v_coords - cy) * z_cam / fy
    
    pc_cam = torch.stack([x_cam, y_cam, z_cam], dim=-1)
    
    # # (n_env, n_cam, H, W)
    mask = (depth > znear) & (depth < zfar * 0.99)
    
    if not world_frame:
        return pc_cam, mask
    
    # Transform to world frame
    pc_cam_h = torch.cat([
        pc_cam,
        torch.ones(n_env, n_cam, H, W, 1, device=device)
    ], dim=-1) # (n_env, n_cam, H, W, 4)
    
    # (n_env, n_cam, H*W, 4)
    pc_cam_h_flat = pc_cam_h.view(n_env, n_cam, H * W, 4)
     
    pc_world_h_flat = pc_cam_h_flat @ pose_opencv_to_world.transpose(-2, -1)
    
    pc_world = pc_world_h_flat[..., :3].view(n_env, n_cam, H, W, 3)
    return pc_world, mask

# Usage of the Function
rgba, depth, _, _ = self.scene.render_all_cameras(rgb=True, depth=True) # rgb - (n_env, H, W, 4) * n_cam tuple; depth - (n_env, H, W, 1) * n_cam tuple
camera_pos = self.scene.visualizer.camera_pos # (n_env, n_cam, 3)
camera_fov = self.scene.visualizer.camera_fov # (30., 30., 30., 30., 30.) 30*n_cam
camera_quat = self.scene.visualizer.camera_quat # (n_env, n_cam, 4)
        
pc, mask = get_pointcloud_batch(
    depth_tuple=depth,
    camera_pos=camera_pos,
    camera_quat=camera_quat,
    camera_fov=camera_fov,
    znear=0.05,
    zfar=100.0,
    world_frame=True
)