initial support for gaussian particles

CMakeLists.txt

@@ -40,10 +40,14 @@ nanobind_add_module(polyscope_bindings
   src/cpp/volume_grid.cpp
   src/cpp/sparse_volume_grid.cpp
   src/cpp/camera_view.cpp
+  src/cpp/gaussian_particles.cpp
   src/cpp/floating_quantities.cpp
   src/cpp/implicit_helpers.cpp
   src/cpp/managed_buffer.cpp
 
+  src/cpp/gaussian_particles_structure.cpp
+  src/cpp/gaussian_particles_structure.h
+
   src/cpp/utils.h
 
   # ImGui related things

src/cpp/core.cpp

@@ -32,6 +32,7 @@ void bind_volume_mesh(nb::module_& m);
 void bind_volume_grid(nb::module_& m);
 void bind_sparse_volume_grid(nb::module_& m);
 void bind_camera_view(nb::module_& m);
+void bind_gaussian_particles(nb::module_& m);
 void bind_floating_quantities(nb::module_& m);
 void bind_implicit_helpers(nb::module_& m);
 void bind_managed_buffer(nb::module_& m);
@@ -59,6 +60,8 @@ NB_MODULE(polyscope_bindings, m) {
         ps::state::userCallback = nullptr;
         if (ps::render::engine != nullptr) {
           ps::shutdown(true);
+        } else {
+          ps::removeEverything();
         }
   }));
 
@@ -87,6 +90,7 @@ NB_MODULE(polyscope_bindings, m) {
   m.def("window_requests_close", &ps::windowRequestsClose);
   m.def("frame_tick", &ps::frameTick);
   m.def("shutdown", &ps::shutdown);
+  m.def("remove_everything", &ps::removeEverything);
 
   // === Render engine related things
   m.def("get_render_engine_backend_name", &ps::render::getRenderEngineBackendName);
@@ -155,6 +159,7 @@ NB_MODULE(polyscope_bindings, m) {
   m.def("get_length_scale", []() { return ps::state::lengthScale; });
   m.def("set_bounding_box", [](glm::vec3 low, glm::vec3 high) { ps::state::boundingBox = std::tuple<glm::vec3, glm::vec3>(low, high); });
   m.def("get_bounding_box", []() { return ps::state::boundingBox; });
+  m.def("update_scene_extents", &ps::updateStructureExtents);
 
   // === Camera controls & View
   m.def("set_navigation_style", [](ps::view::NavigateStyle x) { ps::view::style = x; });
@@ -174,6 +179,7 @@ NB_MODULE(polyscope_bindings, m) {
       ps::view::lookAt(location, target, upDir, flyTo); 
   });
   m.def("set_view_projection_mode", [](ps::ProjectionMode x) { ps::view::projectionMode = x; });
+  m.def("get_view_projection_mode", []() { return ps::view::projectionMode; });
   m.def("get_view_camera_parameters", &ps::view::getCameraParametersForCurrentView);
   m.def("set_view_camera_parameters", &ps::view::setViewToCamera);
   m.def("set_camera_view_matrix", [](Eigen::Matrix4f mat) { ps::view::setCameraViewMatrix(eigen2glm(mat)); });
@@ -730,6 +736,7 @@ NB_MODULE(polyscope_bindings, m) {
   bind_volume_grid(m);
   bind_sparse_volume_grid(m);
   bind_camera_view(m);
+  bind_gaussian_particles(m);
   bind_managed_buffer(m);
   bind_imgui(m);
   bind_implot(m);

src/cpp/gaussian_particles.cpp

@@ -0,0 +1,51 @@
+#include "utils.h"
+
+#include <nanobind/eigen/dense.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/string.h>
+
+#include "gaussian_particles_structure.h"
+
+// === Bindings
+
+// clang-format off
+void bind_gaussian_particles(nb::module_& m) {
+
+  // == Helper classes
+  nb::class_<ps::GaussianParticlesPickResult>(m, "GaussianParticlesPickResult")
+   .def(nb::init<>())
+   .def_ro("index", &ps::GaussianParticlesPickResult::index)
+  ;
+
+  // == Main class
+  bindStructure<ps::GaussianParticles>(m, "GaussianParticles")
+
+    // basics
+
+    // options
+
+    // picking
+
+    // quantities
+
+
+    // 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)
+    .def("set_subsample_factor", &ps::GaussianParticles::setSubsampleFactor)
+
+    ;
+
+  // Static adders and getters
+  m.def("register_gaussian_particles", &ps::registerGaussianParticles, 
+      nb::arg("name"), nb::arg("draw_func"), nb::arg("extents_callback"), nb::arg("deletion_callback"), "Register gaussian particles", 
+      nb::rv_policy::reference);
+  m.def("remove_gaussian_particles", &polyscope::removeGaussianParticles, "Remove gaussian particles by name");
+  m.def("get_gaussian_particles", &polyscope::getGaussianParticles, "Get gaussian particles by name", nb::rv_policy::reference);
+  m.def("has_gaussian_particles", &polyscope::hasGaussianParticles, "Check for gaussian particles by name");
+
+}
+// clang-format on

src/cpp/gaussian_particles_structure.cpp

@@ -0,0 +1,169 @@
+// Copyright 2017-2023, Nicholas Sharp and the Polyscope contributors. https://polyscope.run
+
+#include "gaussian_particles_structure.h"
+
+#include "polyscope/affine_remapper.h"
+#include "polyscope/color_management.h"
+#include "polyscope/persistent_value.h"
+#include "polyscope/pick.h"
+#include "polyscope/polyscope.h"
+#include "polyscope/render/engine.h"
+#include "polyscope/render/managed_buffer.h"
+#include "polyscope/scaled_value.h"
+#include "polyscope/standardize_data_array.h"
+#include "polyscope/structure.h"
+
+namespace polyscope {
+
+// Initialize statics
+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_), 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.
+}
+
+GaussianParticles::~GaussianParticles() { deletionCallback(); }
+
+void GaussianParticles::buildCustomUI() {
+  ensureImagebuffersAllocated(); // doing this here ensures we re-render after resizing
+
+  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) {}
+
+// Standard structure overrides
+void GaussianParticles::draw() {}
+void GaussianParticles::drawDelayed() {
+  if (!isEnabled()) return;
+
+
+  ensureImagebuffersAllocated();
+
+  if (!internal::renderPassIsRedraw) {
+    // if doing multiple passes from the same view, no need to do this multiple times
+    drawCallback();
+  }
+
+  if (!imageToScreenProgram) {
+    prepareImageToScreenProgram();
+  }
+  setImageToScreenUniforms();
+  imageToScreenProgram->draw();
+}
+
+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() {}
+
+std::tuple<int32_t, int32_t> GaussianParticles::getRenderDims() {
+  ensureImagebuffersAllocated();
+  return std::make_tuple(currImageWidth, currImageHeight);
+}
+
+void GaussianParticles::ensureImagebuffersAllocated() {
+  int32_t newImageWidth = view::bufferWidth / subsampleFactor.get();
+  int32_t newImageHeight = view::bufferHeight / subsampleFactor.get();
+
+  if (newImageHeight == currImageHeight && newImageWidth == currImageWidth) {
+    return;
+  }
+
+  currImageHeight = newImageHeight;
+  currImageWidth = newImageWidth;
+
+  depths.reset();
+  colors.reset();
+  depths = std::make_unique<render::ManagedBuffer<float>>(this, uniquePrefix() + "depths", depthsData);
+  colors = std::make_unique<render::ManagedBuffer<glm::vec4>>(this, uniquePrefix() + "colors", colorsData);
+
+  // re-allocate buffers
+  depths->setTextureSize(currImageWidth, currImageHeight);
+  depths->ensureHostBufferAllocated();
+  depths->data = std::vector<float>(currImageWidth * currImageHeight, 0.0f);
+  depths->markHostBufferUpdated();
+
+  colors->setTextureSize(currImageWidth, currImageHeight);
+  colors->ensureHostBufferAllocated();
+  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();
+}
+
+void GaussianParticles::setImageToScreenUniforms() {
+  setStructureUniforms(*imageToScreenProgram);
+
+  glm::mat4 P = view::getCameraPerspectiveMatrix();
+  glm::mat4 Pinv = glm::inverse(P);
+
+  imageToScreenProgram->setUniform("u_projMatrix", glm::value_ptr(P));
+  imageToScreenProgram->setUniform("u_invProjMatrix", glm::value_ptr(Pinv));
+  imageToScreenProgram->setUniform("u_viewport", render::engine->getCurrentViewport());
+  imageToScreenProgram->setUniform("u_textureTransparency", transparency.get());
+  render::engine->setTonemapUniforms(*imageToScreenProgram);
+  if (imageToScreenProgram->hasUniform("u_transparency")) {
+    imageToScreenProgram->setUniform("u_transparency", 1.0f);
+  }
+}
+
+void GaussianParticles::prepareImageToScreenProgram() {
+
+  // NOTE: we use INVERSE_TONEMAP to avoid tonemapping the content, but in the presence of transparency this setup
+  // cannot exactly preserve the result, since the inversion is applied before compositing but finaltonemapping is
+  // applied after compositing.
+
+  // Create the sourceProgram
+  // clang-format off
+  std::vector<std::string> rules = addStructureRules({
+    getImageOriginRule(ImageOrigin::UpperLeft),
+    "TEXTURE_SHADE_COLORALPHA", "INVERSE_TONEMAP"
+  });
+  rules = removeRule(rules, "GENERATE_VIEW_POS");
+
+  imageToScreenProgram = render::engine->requestShader("TEXTURE_DRAW_RAW_RENDERIMAGE_PLAIN", rules);
+  // clang-format on
+
+  imageToScreenProgram->setAttribute("a_position", render::engine->screenTrianglesCoords());
+  imageToScreenProgram->setTextureFromBuffer("t_depth", depths->getRenderTextureBuffer().get());
+  imageToScreenProgram->setTextureFromBuffer("t_color", colors->getRenderTextureBuffer().get());
+}
+
+} // namespace polyscope