initial support for gaussian particles

CMakeLists.txt

@@ -39,10 +39,14 @@ nanobind_add_module(polyscope_bindings
   src/cpp/volume_mesh.cpp
   src/cpp/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

@@ -31,6 +31,7 @@ void bind_curve_network(nb::module_& m);
 void bind_volume_mesh(nb::module_& m);
 void bind_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);
@@ -58,6 +59,8 @@ NB_MODULE(polyscope_bindings, m) {
         ps::state::userCallback = nullptr;
         if (ps::render::engine != nullptr) {
           ps::shutdown(true);
+        } else {
+          ps::removeEverything();
         }
   }));
 
@@ -86,6 +89,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);
@@ -154,6 +158,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; });
@@ -173,6 +178,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)); });
@@ -718,6 +724,7 @@ NB_MODULE(polyscope_bindings, m) {
   bind_volume_mesh(m);
   bind_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,49 @@
+#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("update_object_space_bounds", &ps::GaussianParticles::updateObjectSpaceBounds)
+
+    ;
+
+  // 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,144 @@
+// 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_) {
+
+  // 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");
+}
+void GaussianParticles::buildCustomOptionsUI() {}
+void GaussianParticles::buildPickUI(const PickResult& result) {}
+
+// Standard structure overrides
+void GaussianParticles::draw() {}
+void GaussianParticles::drawDelayed() {
+  if (!isEnabled()) return;
+
+
+  ensureImagebuffersAllocated();
+
+  drawCallback();
+
+  if (!imageToScreenProgram) {
+    prepareImageToScreenProgram();
+  }
+  setImageToScreenUniforms();
+  imageToScreenProgram->draw();
+}
+
+void GaussianParticles::drawPick() {}
+void GaussianParticles::drawPickDelayed() {}
+void GaussianParticles::updateObjectSpaceBounds() { extentsCallback(); }
+
+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();
+}
+
+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;
+  int32_t newImageHeight = view::bufferHeight / subsampleFactor;
+
+  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();
+
+  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

src/cpp/gaussian_particles_structure.h

@@ -0,0 +1,132 @@
+// Copyright 2017-2023, Nicholas Sharp and the Polyscope contributors. https://polyscope.run
+
+#pragma once
+
+#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 {
+
+// Forward declare gaussian particles
+class GaussianParticles;
+
+// Forward declare quantity types
+// class GaussianParticlesScalarQuantity;
+
+
+struct GaussianParticlesPickResult {
+  int64_t index;
+};
+
+class GaussianParticles : public Structure {
+public:
+  // === Member functions ===
+
+  // Construct a new structure
+  GaussianParticles(std::string name, std::function<void()>& drawCallback, std::function<void()>& extentsCallback,
+                    std::function<void()>& deletionCallback);
+  virtual ~GaussianParticles();
+
+  // === Overrides
+
+  // Build the imgui display
+  virtual void buildCustomUI() override;
+  virtual void buildCustomOptionsUI() override;
+  virtual void buildPickUI(const PickResult& result) override;
+
+  // Standard structure overrides
+  virtual void draw() override;
+  virtual void drawDelayed() override;
+  virtual void drawPick() override;
+  virtual void drawPickDelayed() override;
+  virtual void updateObjectSpaceBounds() override;
+  virtual std::string typeName() override;
+  virtual void refresh() override;
+
+  // Misc data
+  static const std::string structureTypeName;
+  std::tuple<int32_t, int32_t> getRenderDims();
+
+  int32_t getSubsampleFactor();
+  void setSubsampleFactor(int32_t newVal);
+
+  // 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);
+
+
+private:
+  std::function<void()> drawCallback;
+  std::function<void()> extentsCallback;
+  std::function<void()> deletionCallback;
+
+  int32_t currImageWidth = -1;
+  int32_t currImageHeight = -1;
+  int32_t subsampleFactor = 1;
+
+  // === Render data
+  std::vector<float> depthsData;
+  std::unique_ptr<render::ManagedBuffer<float>> depths;
+  std::vector<glm::vec4> colorsData;
+  std::unique_ptr<render::ManagedBuffer<glm::vec4>> colors;
+  std::shared_ptr<render::ShaderProgram> imageToScreenProgram;
+
+
+  // === Helpers
+  void ensureImagebuffersAllocated();
+  void setImageToScreenUniforms();
+  void prepareImageToScreenProgram();
+};
+
+// Shorthand to add a gaussian particles to polyscope
+GaussianParticles* registerGaussianParticles(std::string name, std::function<void()>& drawCallback,
+                                             std::function<void()>& extentsCallback,
+                                             std::function<void()>& deletionCallback);
+
+// Shorthand to get a gaussian particles from polyscope
+inline GaussianParticles* getGaussianParticles(std::string name = "");
+inline bool hasGaussianParticles(std::string name = "");
+inline void removeGaussianParticles(std::string name = "", bool errorIfAbsent = false);
+
+
+} // namespace polyscope
+
+
+// implementations of the inline funcitons
+// (this setup is silly from a C++ standpoint, we do it to match other structures)
+
+namespace polyscope {
+
+inline GaussianParticles* registerGaussianParticles(std::string name, std::function<void()>& drawCallback,
+                                                    std::function<void()>& extentsCallback,
+                                                    std::function<void()>& deletionCallback) {
+
+  checkInitialized();
+
+  GaussianParticles* s = new GaussianParticles(name, drawCallback, extentsCallback, deletionCallback);
+  bool success = registerStructure(s);
+  if (!success) {
+    safeDelete(s);
+  }
+  return s;
+}
+
+
+// Shorthand to get a gaussian particles from polyscope
+inline GaussianParticles* getGaussianParticles(std::string name) {
+  return dynamic_cast<GaussianParticles*>(getStructure(GaussianParticles::structureTypeName, name));
+}
+inline bool hasGaussianParticles(std::string name) { return hasStructure(GaussianParticles::structureTypeName, name); }
+inline void removeGaussianParticles(std::string name, bool errorIfAbsent) {
+  removeStructure(GaussianParticles::structureTypeName, name, errorIfAbsent);
+}
+
+} // namespace polyscope