From eb9bcb577d5bc0249eba05f6d11a3b4e53ef8a3d Mon Sep 17 00:00:00 2001
From: "J.Y" <132313008+jb-ye@users.noreply.github.com>
Date: Tue, 13 Feb 2024 14:43:04 -0500
Subject: [PATCH] Implement new rasterization mode using opacity compensation
 factor (#2888)

Implement new rasterization mode for tiny gaussians
---
 nerfstudio/models/splatfacto.py | 25 +++++++++++++++++++++++--
 pyproject.toml                  |  2 +-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py
index faaa2150c..90bdccd70 100644
--- a/nerfstudio/models/splatfacto.py
+++ b/nerfstudio/models/splatfacto.py
@@ -157,6 +157,16 @@ class SplatfactoModelConfig(ModelConfig):
     """
     output_depth_during_training: bool = False
     """If True, output depth during training. Otherwise, only output depth during evaluation."""
+    rasterize_mode: Literal["classic", "antialiased"] = "classic"
+    """
+    Classic mode of rendering will use the EWA volume splatting with a [0.3, 0.3] screen space blurring kernel. This
+    approach is however not suitable to render tiny gaussians at higher or lower resolution than the captured, which
+    results "aliasing-like" artifacts. The antialiased mode overcomes this limitation by calculating compensation factors
+    and apply them to the opacities of gaussians to preserve the total integrated density of splats.
+
+    However, PLY exported with antialiased rasterize mode is not compatible with classic mode. Thus many web viewers that
+    were implemented for classic mode can not render antialiased mode PLY properly without modifications.
+    """
 
 
 class SplatfactoModel(Model):
@@ -721,7 +731,7 @@ class SplatfactoModel(Model):
 
         colors_crop = torch.cat((features_dc_crop[:, None, :], features_rest_crop), dim=1)
 
-        self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = project_gaussians(  # type: ignore
+        self.xys, depths, self.radii, conics, comp, num_tiles_hit, cov3d = project_gaussians(  # type: ignore
             means_crop,
             torch.exp(scales_crop),
             1,
@@ -736,6 +746,7 @@ class SplatfactoModel(Model):
             W,
             tile_bounds,
         )  # type: ignore
+
         if (self.radii).sum() == 0:
             rgb = background.repeat(int(camera.height.item()), int(camera.width.item()), 1)
             depth = background.new_ones(*rgb.shape[:2], 1) * 10
@@ -758,6 +769,16 @@ class SplatfactoModel(Model):
         # rescale the camera back to original dimensions
         camera.rescale_output_resolution(camera_downscale)
         assert (num_tiles_hit > 0).any()  # type: ignore
+
+        # apply the compensation of screen space blurring to gaussians
+        opacities = None
+        if self.config.rasterize_mode == "antialiased":
+            opacities = torch.sigmoid(opacities_crop) * comp[:, None]
+        elif self.config.rasterize_mode == "classic":
+            opacities = torch.sigmoid(opacities_crop)
+        else:
+            raise ValueError("Unknown rasterize_mode: %s", self.config.rasterize_mode)
+
         rgb, alpha = rasterize_gaussians(  # type: ignore
             self.xys,
             depths,
@@ -765,7 +786,7 @@ class SplatfactoModel(Model):
             conics,
             num_tiles_hit,  # type: ignore
             rgbs,
-            torch.sigmoid(opacities_crop),
+            opacities,
             H,
             W,
             background=background,
diff --git a/pyproject.toml b/pyproject.toml
index cde1003e2..e2a731f64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,7 +63,7 @@ dependencies = [
     "xatlas",
     "trimesh>=3.20.2",
     "timm==0.6.7",
-    "gsplat==0.1.3",
+    "gsplat>=0.1.5",
     "pytorch-msssim",
     "pathos",
     "packaging"