cleanup to options and some bugfixes

Author: nmwsharp

deps/polyscope

@@ -33,6 +33,7 @@ void bind_gaussian_particles(nb::module_& m) {
     // internal
     .def("get_render_dims", &ps::GaussianParticles::getRenderDims)
     .def("set_extents", &ps::GaussianParticles::setExtents)
+    .def("set_num_particles", &ps::GaussianParticles::setNumParticles)
     .def("update_object_space_bounds", &ps::GaussianParticles::updateObjectSpaceBounds)
 
     ;

src/cpp/gaussian_particles.cpp

@@ -21,18 +21,24 @@ const std::string GaussianParticles::structureTypeName = "Gaussian Particles";
 GaussianParticles::GaussianParticles(std::string name_, std::function<void()>& drawCallback_,
                                      std::function<void()>& extentsCallback_, std::function<void()>& deletionCallback_)
     : Structure(name_, structureTypeName), drawCallback(drawCallback_), extentsCallback(extentsCallback_),
-      deletionCallback(deletionCallback_) {
+      deletionCallback(deletionCallback_), subsampleFactor(uniquePrefix() + "subsampleFactor", 1) {
 
-  // note: unlike other structures this does not call updateObjectSpaceBounds() here, to avoid a circular problem with
-  // the external class. we call it manually right after creation there.
+  // Note: unlike other structures this does not call updateObjectSpaceBounds() here, to avoid a circular problem with
+  // the external class. We call it manually right after creation there.
 }
 
 GaussianParticles::~GaussianParticles() { deletionCallback(); }
 
 void GaussianParticles::buildCustomUI() {
   ensureImagebuffersAllocated(); // doing this here ensures we re-render after resizing
 
-  ImGui::Text("# particles: -1");
+  ImGui::Text("%d Gaussians", numParticles);
+
+  ImGui::PushItemWidth(100);
+  if (ImGui::InputInt("render subsample", &(subsampleFactor.get()), 1, 10)) {
+    setSubsampleFactor(subsampleFactor.get());
+  }
+  ImGui::PopItemWidth();
 }
 void GaussianParticles::buildCustomOptionsUI() {}
 void GaussianParticles::buildPickUI(const PickResult& result) {}
@@ -45,7 +51,10 @@ void GaussianParticles::drawDelayed() {
 
   ensureImagebuffersAllocated();
 
-  drawCallback();
+  if (!internal::renderPassIsRedraw) {
+    // if doing multiple passes from the same view, no need to do this multiple times
+    drawCallback();
+  }
 
   if (!imageToScreenProgram) {
     prepareImageToScreenProgram();
@@ -58,13 +67,23 @@ void GaussianParticles::drawPick() {}
 void GaussianParticles::drawPickDelayed() {}
 void GaussianParticles::updateObjectSpaceBounds() { extentsCallback(); }
 
+void GaussianParticles::setSubsampleFactor(int32_t newVal) {
+  newVal = std::max(1, newVal);
+  subsampleFactor = newVal;
+  ensureImagebuffersAllocated();
+}
+
+int32_t GaussianParticles::getSubsampleFactor() { return subsampleFactor.get(); }
+
 void GaussianParticles::setExtents(glm::vec3 bbox_min, glm::vec3 bbox_max) {
   std::get<0>(objectSpaceBoundingBox) = bbox_min;
   std::get<1>(objectSpaceBoundingBox) = bbox_max;
   objectSpaceLengthScale = glm::length(bbox_max - bbox_min);
   requestRedraw();
 }
 
+void GaussianParticles::setNumParticles(int32_t numParticles_) { numParticles = numParticles_; }
+
 std::string GaussianParticles::typeName() { return structureTypeName; }
 void GaussianParticles::refresh() {}
 
@@ -74,8 +93,8 @@ std::tuple<int32_t, int32_t> GaussianParticles::getRenderDims() {
 }
 
 void GaussianParticles::ensureImagebuffersAllocated() {
-  int32_t newImageWidth = view::bufferWidth / subsampleFactor;
-  int32_t newImageHeight = view::bufferHeight / subsampleFactor;
+  int32_t newImageWidth = view::bufferWidth / subsampleFactor.get();
+  int32_t newImageHeight = view::bufferHeight / subsampleFactor.get();
 
   if (newImageHeight == currImageHeight && newImageWidth == currImageWidth) {
     return;
@@ -100,6 +119,12 @@ void GaussianParticles::ensureImagebuffersAllocated() {
   colors->data = std::vector<glm::vec4>(currImageWidth * currImageHeight, glm::vec4(0.0f));
   colors->markHostBufferUpdated();
 
+  if (imageToScreenProgram) {
+    imageToScreenProgram->setTextureFromBuffer("t_depth", depths->getRenderTextureBuffer().get());
+    imageToScreenProgram->setTextureFromBuffer("t_color", colors->getRenderTextureBuffer().get());
+  }
+
+
   requestRedraw();
 }
 

src/cpp/gaussian_particles_structure.cpp

@@ -61,6 +61,7 @@ class GaussianParticles : public Structure {
   // Manually set the extents of the structure (note: does _not_ automatically update the scene, call refresh() after if
   // you need that.)
   void setExtents(glm::vec3 bbox_min, glm::vec3 bbox_max);
+  void setNumParticles(int32_t numParticles);  
 
 
 private:
@@ -70,7 +71,10 @@ class GaussianParticles : public Structure {
 
   int32_t currImageWidth = -1;
   int32_t currImageHeight = -1;
-  int32_t subsampleFactor = 1;
+  int32_t currSubsampleFactor = -1;
+  int32_t numParticles = -1;
+
+  PersistentValue<int32_t> subsampleFactor;
 
   // === Render data
   std::vector<float> depthsData;

src/cpp/gaussian_particles_structure.h

@@ -93,6 +93,7 @@ def structure_extents_callback(self):
         max_bound = np.percentile(centers, 75, axis=0)
 
         self.bound_instance.set_extents(glm3(min_bound), glm3(max_bound))
+        self.bound_instance.set_num_particles(centers.shape[0])
 
     def structure_deletion_callback(self):
         """
@@ -106,9 +107,8 @@ def n_particles(self) -> int:
         return self.bound_instance.n_particles()
 
     def draw(self):
-        import gsplat, torch
-
-        # print("drawing gaussian particles!")
+        import torch
+        import gsplat
 
         color_buffer = self.get_buffer("colors")
         depth_buffer = self.get_buffer("depths")
@@ -118,6 +118,7 @@ def draw(self):
 
         with torch.no_grad():
 
+            # prep camera matrices for the current Polyscope view
             render_w, render_h = self.bound_instance.get_render_dims()
             view_params = get_view_camera_parameters()
             T_camera_convent = np.array([
@@ -126,7 +127,6 @@ def draw(self):
                 [0., 0., -1., 0.],
                 [0., 0., 0., 1.],
             ], dtype=np.float32)
-
             view_mat = torch.tensor(T_camera_convent @ view_params.get_view_mat(), device=self.device)
             fov_height = np.deg2rad(view_params.get_fov_vertical_deg())
             fov_width = 2.0 * np.arctan(np.tan(fov_height / 2.0) * view_params.get_aspect())
@@ -145,7 +145,6 @@ def draw(self):
                 width=render_w,
                 height=render_h,
                 render_mode="RGB+ED",
-                sh_degree=0
             )
 
             # unbatchify

src/polyscope/gaussian_particles.py

@@ -21,16 +21,7 @@
 import argparse
 import numpy as np
 from plyfile import PlyData
-from arrgh import arrgh
 
-'''
-Download the mipnerf360 dataset:
-
-mkdir -p data/mipnerf360
-cd data/mipnerf360
-wget http://storage.googleapis.com/gresearch/refraw360/360_v2.zip
-unzip 360_v2.zip
-'''
 
 def load_gaussians_from_ply(path_ply, device='cuda'):
 
@@ -91,33 +82,28 @@ def main():
     args = parser.parse_args()
 
 
-    ps.init() 
-    ps.set_ground_plane_mode("shadow_only")
+    ps.init()
+    ps.set_ground_plane_mode("none")
 
-    def callback():
-        pass
-    ps.set_user_callback(callback)
-   
-    
     # Load the file
     centers, features_dc, opacity, scaling, rotation = load_gaussians_from_ply(args.gaussian_ply)
     print(f"Loaded {centers.shape[0]} gaussian particles")
 
-
-    arrgh(centers, features_dc, opacity, scaling, rotation)
-
-    ps.register_gaussian_particles("gaussians", 
-                                   means=centers.unsqueeze(0), 
-                                   colors=features_dc.unsqueeze(0), 
-                                   opacities=opacity.unsqueeze(0), 
-                                   scales=scaling.unsqueeze(0), 
-                                   quats=rotation.unsqueeze(0)
+    # Register Gaussians
+    ps.register_gaussian_particles("gaussians",
+                                   # the arguments below are passed directly
+                                   # to gsplat.rasterization(...). See the docs there for the meaning
+                                   # and other parameters you can pass.
+                                   means=centers.unsqueeze(0),
+                                   colors=features_dc.unsqueeze(0),
+                                   opacities=opacity.unsqueeze(0),
+                                   scales=scaling.unsqueeze(0),
+                                   quats=rotation.unsqueeze(0),
+                                   sh_degree=0,
                                 )
 
     ps.show()
 
-    # ps.remove_gaussian_particles("gaussians")
-
 
 if __name__ == '__main__':
     main()