2DGS

.gitignore

@@ -121,5 +121,8 @@ venv.bak/
 compile_commands.json
 *.dump
 
-data
-results
+/data
+/result
+/results
+/renders
+*.png

examples/2d_example.py

@@ -0,0 +1,124 @@
+import math
+import os
+from pathlib import Path
+from typing import Optional, Literal
+
+import torch
+import numpy as np
+import tyro
+import matplotlib
+from PIL import Image
+
+from gsplat import rasterization, rasterization_2dgs
+
+class SimpleTrainer:
+
+    def __init__(
+        self,
+        num_points_per_axis: int = 8,
+    ):
+        self.device = torch.device("cuda:0")
+        self.num_points_per_axis = num_points_per_axis
+
+        fov_x = math.pi / 2.0
+        self.H, self.W = 256, 256
+        self.focal = 0.5 * float(self.W) / math.tan(0.5 * fov_x)
+        self.img_size = torch.tensor([self.W, self.H, 1], device=self.device)
+
+        self._init_gaussians()
+
+    def _init_gaussians(self):
+        length = 0.4
+        x = np.linspace(-1, 1, self.num_points_per_axis)
+        y = np.linspace(-1, 1, self.num_points_per_axis)
+        x, y = np.meshgrid(x, y)
+        means3D = torch.from_numpy(np.stack([x, y, np.ones_like(x)], axis=-1).reshape(-1, 3)).cuda().float()
+        quats = torch.zeros(1, 4).repeat(len(means3D), 1).cuda()
+        quats[..., 0] = 1.
+        scale = 0.6 / (self.num_points_per_axis - 1)
+        scales = torch.zeros(1, 3).repeat(len(means3D), 1).fill_(scale).cuda()
+        num_points = self.num_points_per_axis ** 2
+        colors = matplotlib.colormaps['Accent'](np.random.randint(1, num_points, num_points) / num_points)[..., :3]
+        colors = torch.from_numpy(colors).cuda()
+        opacity = torch.ones_like(means3D[:, 0])
+
+        self.viewmat = torch.tensor(
+            [[-8.6086e-01,  3.7950e-01, -3.3896e-01,  6.7791e-01],
+            [ 5.0884e-01,  6.4205e-01, -5.7346e-01,  1.1469e+00],
+            [ 1.0934e-08, -6.6614e-01, -7.4583e-01,  1.4917e+00],
+            [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00]],
+            device=self.device,
+        )
+
+        self.means = means3D.float()
+        self.scales = scales.float()
+        self.rgbs = colors.float()
+        self.opacities = opacity.float()
+        self.quats = quats.float()
+
+    def render(
+        self,
+        model_type: Literal["3dgs", "2dgs"] = "3dgs",
+    ):
+        frames = []
+        K = torch.tensor(
+            [
+                [self.focal, 0, self.W / 2],
+                [0, self.focal, self.H / 2],
+                [0, 0, 1],
+            ],
+            device=self.device,
+        )
+        if model_type == "3dgs":
+            renders, _ = rasterization(
+                self.means,
+                self.quats / self.quats.norm(dim=-1, keepdim=True),
+                self.scales,
+                self.opacities,
+                self.rgbs,
+                self.viewmat[None],
+                K[None],
+                self.W,
+                self.H,
+                packed=False,
+            )
+
+        elif model_type == "2dgs":
+            renders, _ = rasterization_2dgs(
+                self.means,
+                self.quats / self.quats.norm(dim=-1, keepdim=True),
+                self.scales,
+                self.opacities,
+                self.rgbs,
+                self.viewmat[None],
+                K[None],
+                self.W,
+                self.H,
+                packed=False,
+            )
+        else:
+            raise NotImplementedError("Model not implemented")
+        out_img = renders[0].squeeze(0)
+        torch.cuda.synchronize()
+
+        frame = (out_img.detach().cpu().numpy() * 255).astype(np.uint8)
+        frame_img = Image.fromarray(frame)
+        out_dir = os.path.join(os.getcwd(), "renders")
+        os.makedirs(out_dir, exist_ok=True)
+        frame_img.save(f"{out_dir}/{model_type}.png")
+
+def main(
+    height: int = 256,
+    width: int = 256,
+    num_points_per_axis: int = 8,
+    save_imgs: bool = True,
+    img_path: Optional[Path] = None,
+    model_type: Literal["3dgs", "2dgs"] = "3dgs",
+) -> None:
+    trainer = SimpleTrainer(num_points_per_axis=num_points_per_axis)
+    trainer.render(
+        model_type=model_type,
+    )
+
+if __name__ == "__main__":
+    tyro.cli(main)

examples/image_fitting.py

@@ -2,16 +2,18 @@
 import os
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Literal
 
 import numpy as np
 import torch
 import tyro
 from PIL import Image
 from torch import Tensor, optim
 
-from gsplat import rasterization
+from gsplat import rasterization, rasterization_2dgs
 
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+os.environ['TORCH_USE_CUDA_DSA'] = '1'
 
 class SimpleTrainer:
     """Trains random gaussians to fit an image."""
@@ -79,6 +81,7 @@ def train(
         iterations: int = 1000,
         lr: float = 0.01,
         save_imgs: bool = False,
+        model_type: Literal["3dgs", "2dgs"] = "3dgs",
     ):
         optimizer = optim.Adam(
             [self.rgbs, self.means, self.scales, self.opacities, self.quats], lr
@@ -94,9 +97,27 @@ def train(
             ],
             device=self.device,
         )
+
+        if model_type == "3dgs":
+            rasterize_fnc = rasterization
+        elif model_type == "2dgs":
+            rasterize_fnc = rasterization_2dgs
+
         for iter in range(iterations):
             start = time.time()
-            renders, _, _ = rasterization(
+            # renders, _, _ = rasterization(
+            #     self.means,
+            #     self.quats / self.quats.norm(dim=-1, keepdim=True),
+            #     self.scales,
+            #     torch.sigmoid(self.opacities),
+            #     torch.sigmoid(self.rgbs),
+            #     self.viewmat[None],
+            #     K[None],
+            #     self.W,
+            #     self.H,
+            #     packed=False,
+            # )
+            renders, _ = rasterize_fnc(
                 self.means,
                 self.quats / self.quats.norm(dim=-1, keepdim=True),
                 self.scales,
@@ -108,7 +129,7 @@ def train(
                 self.H,
                 packed=False,
             )
-            out_img = renders[0]
+            out_img = renders[0].squeeze(0)
             torch.cuda.synchronize()
             times[0] += time.time() - start
             loss = mse_loss(out_img, self.gt_image)
@@ -122,6 +143,8 @@ def train(
 
             if save_imgs and iter % 5 == 0:
                 frames.append((out_img.detach().cpu().numpy() * 255).astype(np.uint8))
+            # break
+
         if save_imgs:
             # save them as a gif with PIL
             frames = [Image.fromarray(frame) for frame in frames]
@@ -158,6 +181,7 @@ def main(
     img_path: Optional[Path] = None,
     iterations: int = 1000,
     lr: float = 0.01,
+    model_type: Literal["3dgs", "2dgs"] = "3dgs",
 ) -> None:
     if img_path:
         gt_image = image_path_to_tensor(img_path)
@@ -172,6 +196,7 @@ def main(
         iterations=iterations,
         lr=lr,
         save_imgs=save_imgs,
+        model_type=model_type,
     )
 
 

examples/requirements.txt

@@ -16,3 +16,4 @@ opencv-python
 tyro
 Pillow
 tensorboard
+matplotlib