Fix 2DGS for Correct Normal and Median Depth Calculation (#430)

* Fix: Add missing camera offsets for render outputs in 2DGS kernel

* Fix: Correct render_normals transformation

* black formatting

gsplat/cuda/csrc/rasterize_to_pixels_2dgs_fwd.cu

@@ -72,6 +72,10 @@ __global__ void rasterize_to_pixels_fwd_2dgs_kernel(
     render_colors += camera_id * image_height * image_width * COLOR_DIM;  // get the global offset of the pixel w.r.t the camera
     render_alphas += camera_id * image_height * image_width;  // get the global offset of the pixel w.r.t the camera
     last_ids += camera_id * image_height * image_width;  // get the global offset of the pixel w.r.t the camera
+    render_normals += camera_id * image_height * image_width * 3;
+    render_distort += camera_id * image_height * image_width;
+    render_median += camera_id * image_height * image_width;
+    median_ids += camera_id * image_height * image_width;
 
     // get the global offset of the background and mask
     if (backgrounds != nullptr) {

gsplat/rendering.py

@@ -1306,7 +1306,9 @@ def rasterization_2dgs(
         "gradient_2dgs": densify,  # This holds the gradient used for densification for 2dgs
     }
 
-    render_normals = render_normals @ torch.linalg.inv(viewmats)[0, :3, :3].T
+    render_normals = torch.einsum(
+        "...ij,...hwj->...hwi", torch.linalg.inv(viewmats)[..., :3, :3], render_normals
+    )
 
     return (
         render_colors,