2DGS

.gitignore

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

examples/image_fitting.py

@@ -2,7 +2,7 @@
 import os
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Literal
 
 import numpy as np
 import torch
@@ -79,6 +79,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

gsplat/cuda/csrc/bindings.h

@@ -212,4 +212,59 @@ fully_fused_projection_packed_bwd_tensor(
     const torch::Tensor &v_depths,                      // [nnz]
     const torch::Tensor &v_conics,                      // [nnz, 3]
     const at::optional<torch::Tensor> &v_compensations, // [nnz] optional
-    const bool viewmats_requires_grad, const bool sparse_grad);
+    const bool viewmats_requires_grad, const bool sparse_grad);
+
+//====== 2DGS ======//
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+fully_fused_projection_fwd_2dgs_tensor(
+    const int num_points,
+    const torch::Tensor &means,     // [N, 3]
+    const torch::Tensor &quats,     // [N, 4]
+    const torch::Tensor &scales,    // [N, 3]
+    const torch::Tensor &viewmat,   // [C, 4, 4]
+    const torch::Tensor &Ks,        // [C, 3, 3]      
+    const uint32_t image_width, const uint32_t image_height, const float eps2d,
+    const float near_plane, const float far_plane,
+    const float radius_clip, const bool calc_compensations
+);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, 
+            torch::Tensor, torch::Tensor>
+fully_fused_projection_packed_fwd_2dgs_tensor(
+    const torch::Tensor &means,
+    const torch::Tensor &quats,
+    const torch::Tensor &scales,
+    const torch::Tensor &viewmats,
+    const torch::Tensor &Ks,
+    const uint32_t image_width, const uint32_t image_height, const float eps2d,
+    const float near_plane, const float far_plane, const float radius_clip,
+    const bool calc_compensations
+);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> rasterize_to_pixels_fwd_2dgs_tensor(
+    // Gaussian parameters
+    const torch::Tensor &means2d,
+    const torch::Tensor &colors,
+    const torch::Tensor &ray_transformations,
+    const torch::Tensor &opacities,
+    const at::optional<torch::Tensor> &backgrounds,
+    // image size
+    const uint32_t image_width, const uint32_t image_height, const uint32_t tile_size,
+    // intersections
+    const torch::Tensor &tile_offsets, // [C, tile_height, tile_width]
+    const torch::Tensor &flatten_ids   // [n_isects]
+);
+
+std::tuple<torch::Tensor, torch::Tensor> rasterize_to_indices_in_range_2dgs_tensor(
+    const uint32_t range_start, const uint32_t range_end,   // iteration steps
+    const torch::Tensor transmittances, // [C, image_height, image_width]
+    // Gaussian parameters
+    const torch::Tensor &means2d,   // [C, N, 2]
+    const torch::Tensor &ray_transformations, // [C, N, 3, 3]
+    const torch::Tensor &opacities, // [N]
+    // image size
+    const uint32_t image_width, const uint32_t image_height, const uint32_t tile_size, 
+    // intersections
+    const torch::Tensor &tile_offsets, // [C, tile_height, tile_width]
+    const torch::Tensor &flatten_ids    // [n_isects]
+);