diff --git a/CMakeLists.txt b/CMakeLists.txt index d5060c3..c26623c 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/deps/polyscope b/deps/polyscope index 150d201..3697a3d 160000 --- a/deps/polyscope +++ b/deps/polyscope @@ -1 +1 @@ -Subproject commit 150d201bbe7f39e91a5a8beb008ea04b59bbd7bb +Subproject commit 3697a3d8174e50004584cc3d9876d2d6cb682051 diff --git a/src/cpp/core.cpp b/src/cpp/core.cpp index 3175d51..68cc242 100644 --- a/src/cpp/core.cpp +++ b/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(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); diff --git a/src/cpp/gaussian_particles.cpp b/src/cpp/gaussian_particles.cpp new file mode 100644 index 0000000..b960b4c --- /dev/null +++ b/src/cpp/gaussian_particles.cpp @@ -0,0 +1,49 @@ +#include "utils.h" + +#include +#include +#include +#include + +#include "gaussian_particles_structure.h" + +// === Bindings + +// clang-format off +void bind_gaussian_particles(nb::module_& m) { + + // == Helper classes + nb::class_(m, "GaussianParticlesPickResult") + .def(nb::init<>()) + .def_ro("index", &ps::GaussianParticlesPickResult::index) + ; + + // == Main class + bindStructure(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 \ No newline at end of file diff --git a/src/cpp/gaussian_particles_structure.cpp b/src/cpp/gaussian_particles_structure.cpp new file mode 100644 index 0000000..d3e8010 --- /dev/null +++ b/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& drawCallback_, + std::function& extentsCallback_, std::function& 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 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>(this, uniquePrefix() + "depths", depthsData); + colors = std::make_unique>(this, uniquePrefix() + "colors", colorsData); + + // re-allocate buffers + depths->setTextureSize(currImageWidth, currImageHeight); + depths->ensureHostBufferAllocated(); + depths->data = std::vector(currImageWidth * currImageHeight, 0.0f); + depths->markHostBufferUpdated(); + + colors->setTextureSize(currImageWidth, currImageHeight); + colors->ensureHostBufferAllocated(); + colors->data = std::vector(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 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 diff --git a/src/cpp/gaussian_particles_structure.h b/src/cpp/gaussian_particles_structure.h new file mode 100644 index 0000000..c0c186b --- /dev/null +++ b/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& drawCallback, std::function& extentsCallback, + std::function& 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 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 drawCallback; + std::function extentsCallback; + std::function deletionCallback; + + int32_t currImageWidth = -1; + int32_t currImageHeight = -1; + int32_t subsampleFactor = 1; + + // === Render data + std::vector depthsData; + std::unique_ptr> depths; + std::vector colorsData; + std::unique_ptr> colors; + std::shared_ptr imageToScreenProgram; + + + // === Helpers + void ensureImagebuffersAllocated(); + void setImageToScreenUniforms(); + void prepareImageToScreenProgram(); +}; + +// Shorthand to add a gaussian particles to polyscope +GaussianParticles* registerGaussianParticles(std::string name, std::function& drawCallback, + std::function& extentsCallback, + std::function& 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& drawCallback, + std::function& extentsCallback, + std::function& 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(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 \ No newline at end of file diff --git a/src/polyscope/__init__.py b/src/polyscope/__init__.py index 38361df..4a3cb35 100644 --- a/src/polyscope/__init__.py +++ b/src/polyscope/__init__.py @@ -54,4 +54,5 @@ from polyscope.volume_mesh import * # noqa F403 from polyscope.volume_grid import * # noqa F403 from polyscope.camera_view import * # noqa F403 +from polyscope.gaussian_particles import * # noqa F403 from polyscope.global_floating_quantity_structure import * # noqa F403 diff --git a/src/polyscope/core.py b/src/polyscope/core.py index ecd9249..a217f49 100644 --- a/src/polyscope/core.py +++ b/src/polyscope/core.py @@ -60,6 +60,10 @@ def shutdown(allow_mid_frame_shutdown: bool = False) -> None: psb.shutdown(allow_mid_frame_shutdown) +def remove_everything() -> None: + psb.remove_everything() + + ### Render engine @@ -323,6 +327,11 @@ def get_bounding_box(): return np.array(low.as_tuple()), np.array(high.as_tuple()) +def update_scene_extents(): + """This is usually an internal function, and should not need to be called by end users unless you are manually updating structure bounds after creation.""" + psb.update_scene_extents() + + ### Camera controls @@ -338,6 +347,10 @@ def look_at_dir(camera_location, target, up_dir, fly_to: bool = False) -> None: psb.look_at_dir(glm3(camera_location), glm3(target), glm3(up_dir), fly_to) +def get_view_projection_mode() -> Literal["orthographic", "perspective"]: + return from_enum(psb.get_view_projection_mode()) + + def set_view_projection_mode(s: Literal["orthographic", "perspective"] | str) -> None: psb.set_view_projection_mode(to_enum(psb.ProjectionMode, s)) @@ -1034,35 +1047,35 @@ def generate_camera_ray_corners(self): def glm2(vals: ArrayLike) -> psb.glm_vec2: vals_arr = np.asarray(vals, dtype=np.float32) if vals_arr.shape != (2,): - raise ValueError("vals should be length-2 float array/tuple/etc") + raise ValueError("vals should be length-2 float array/tuple/etc. Got shape " + str(vals_arr.shape)) return psb.glm_vec2(vals_arr[0], vals_arr[1]) def glm2i(vals: ArrayLike) -> psb.glm_ivec2: vals_arr = np.asarray(vals, dtype=np.int32) if vals_arr.shape != (2,): - raise ValueError("vals should be length-2 int array/tuple/etc") + raise ValueError("vals should be length-2 int array/tuple/etc. Got shape " + str(vals_arr.shape)) return psb.glm_ivec2(vals_arr[0], vals_arr[1]) def glm3u(vals: ArrayLike) -> psb.glm_uvec3: vals_arr = np.asarray(vals, dtype=np.uint32) if vals_arr.shape != (3,): - raise ValueError("vals should be length-3 int array/tuple/etc") + raise ValueError("vals should be length-3 int array/tuple/etc. Got shape " + str(vals_arr.shape)) return psb.glm_uvec3(vals_arr[0], vals_arr[1], vals_arr[2]) def glm3(vals: ArrayLike) -> psb.glm_vec3: vals_arr = np.asarray(vals, dtype=np.float32) if vals_arr.shape != (3,): - raise ValueError("vals should be length-3 float array/tuple/etc") + raise ValueError("vals should be length-3 float array/tuple/etc. Got shape " + str(vals_arr.shape)) return psb.glm_vec3(vals_arr[0], vals_arr[1], vals_arr[2]) def glm4(vals: ArrayLike) -> psb.glm_vec4: vals_arr = np.asarray(vals, dtype=np.float32) if vals_arr.shape != (4,): - raise ValueError("vals should be length-4 float array/tuple/etc") + raise ValueError("vals should be length-4 float array/tuple/etc. Got shape " + str(vals_arr.shape)) return psb.glm_vec4(vals_arr[0], vals_arr[1], vals_arr[2], vals_arr[3]) diff --git a/src/polyscope/gaussian_particles.py b/src/polyscope/gaussian_particles.py new file mode 100644 index 0000000..6b89ca5 --- /dev/null +++ b/src/polyscope/gaussian_particles.py @@ -0,0 +1,217 @@ +import sys +from typing import Any, Literal, overload, cast + +import polyscope_bindings as psb +from polyscope.core import glm3, get_view_camera_parameters, get_view_projection_mode, update_scene_extents +from polyscope.enums import to_enum, from_enum +from polyscope.structure import Structure + +import numpy as np +from numpy.typing import NDArray, ArrayLike + +if sys.version_info >= (3, 11): + from typing import Unpack +else: + from typing_extensions import Unpack + + +def check_have_dependencies(): + try: + import gsplat + except ImportError: + raise ImportError("gsplat package is required for Gaussian particles. Try `pip install gsplat`") + + try: + import torch + except ImportError: + raise ImportError( + "pytorch package is required for Gaussian particles. See https://pytorch.org/. It is a requred dependency of gsplat, so you should have gotten it when installing gsplat." + ) + + +class GaussianParticles(Structure): + # This class wraps a _reference_ to the underlying object, whose lifetime is managed by Polyscope + # (although, GaussianParticles is a bit different from other structures in that the underlying data is not owned on the C++ side, it's own here in the Python side) + bound_instance: psb.GaussianParticles + device : str + + @overload + def __init__(self, name: str, points: ArrayLike) -> None: ... + + # End users should not call this constrctor, use register_point_cloud instead + def __init__( + self, + name: str | None = None, + instance: psb.GaussianParticles | None = None, + **gaussian_particles_kwargs: Any + ) -> None: + super().__init__() + + if instance is not None: + # Wrap an existing instance + self.bound_instance = instance + + else: + # Create a new instance + assert name is not None + self.validate_gaussian_kwargs(**gaussian_particles_kwargs) + self.gaussian_particles_kwargs = gaussian_particles_kwargs + + self.bound_instance = psb.register_gaussian_particles(name, self.draw, self.structure_extents_callback, self.structure_deletion_callback) + self.bound_instance.update_object_space_bounds() # call this manually to avoid a circular problem + update_scene_extents() + + self.device = "cuda" + + def check_shape(self, points: NDArray) -> None: + # Helper to validate arrays + + if (len(points.shape) != 2) or (points.shape[1] not in (2, 3)): + raise ValueError("Point cloud positions should have shape (N,3); shape is " + str(points.shape)) + + def validate_gaussian_kwargs(self, **kwargs) -> None: + + if "means" not in kwargs: + raise ValueError("missing required argument 'means'") + + if "colors" not in kwargs: + raise ValueError("missing required argument 'colors'") + + for key in ["viewmats", "Ks", "width", "height", "render_mode"]: + if key in kwargs: + raise ValueError(f"don't pass {key} as a kwarg for gsplat, it is set internally by polyscope to render") + + def structure_extents_callback(self): + """ + Used to update the bounding boxes on the internal Polyscope structure + """ + + centers = self.gaussian_particles_kwargs["means"].cpu().numpy().reshape(-1, 3) + + # 25th percentile bbox for some robustness to outliers + min_bound = np.percentile(centers, 25, axis=0) + max_bound = np.percentile(centers, 75, axis=0) + + self.bound_instance.set_extents(glm3(min_bound), glm3(max_bound)) + + def structure_deletion_callback(self): + """ + Gets called when the underlying C++ structure is deleted, which can happen if the user calls remove_gaussian_particles, or if they call remove_everything or shutdown. + """ + name = self.bound_instance.name + if name in gaussian_particles_instance_cache and gaussian_particles_instance_cache[name] is self: + del gaussian_particles_instance_cache[name] + + def n_particles(self) -> int: + return self.bound_instance.n_particles() + + def draw(self): + import gsplat, torch + + # print("drawing gaussian particles!") + + color_buffer = self.get_buffer("colors") + depth_buffer = self.get_buffer("depths") + + if get_view_projection_mode() != "perspective": + raise NotImplementedError("gaussian particles currently only support perspective projection") + + with torch.no_grad(): + + render_w, render_h = self.bound_instance.get_render_dims() + view_params = get_view_camera_parameters() + T_camera_convent = np.array([ + [1., 0., 0., 0.], + [0., -1., 0., 0.], + [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()) + f_x = render_w / (2.0 * np.tan(fov_width / 2.0)) + f_y = render_h / (2.0 * np.tan(fov_height / 2.0)) + proj_mat = torch.tensor([ + [f_x, 0., render_w / 2.0], + [0., f_y, render_h / 2.0], + [0., 0., 1.], + ], dtype=torch.float32, device=self.device) + + render_color, render_alpha, aux = gsplat.rasterization( + **self.gaussian_particles_kwargs, + viewmats=view_mat[None,None,...], + Ks=proj_mat[None,None,...], + width=render_w, + height=render_h, + render_mode="RGB+ED", + sh_degree=0 + ) + + # unbatchify + render_color = render_color[0,...] + render_alpha = render_alpha[0,...] + + # depth component is extra channel of color + render_depth = render_color[..., 3] + render_color = render_color[..., :3] + render_coloralpha = torch.concatenate((render_color, render_alpha), dim=-1) + + # copy to the image buffer + color_buffer.update_data_from_device(render_coloralpha.contiguous()) + depth_buffer.update_data_from_device(render_depth.contiguous()) + + +# A cache of instances, keyed on their name. +# For all other structures, we just create new wrapper objects on-demand. +# This one is different, because the data is owned on the Python side, by a persistant instance of the Python class. +# However, we still want to offer the same API, and not force users to manually keep track of the object. +# As such, we maintain a cache of the "wrapper" objects. +gaussian_particles_instance_cache: dict[str, GaussianParticles] = {} + + +def register_gaussian_particles( + name: str, + enabled: bool | None = None, + transparency: float | None = None, + **gaussian_particles_kwargs: Any +) -> GaussianParticles: + """Register a new point cloud""" + if not psb.is_initialized(): + raise RuntimeError("Polyscope has not been initialized") + + check_have_dependencies() + + p = GaussianParticles(name, **gaussian_particles_kwargs) + gaussian_particles_instance_cache[name] = p + + # == Apply options + if enabled is not None: + p.set_enabled(enabled) + if transparency is not None: + p.set_transparency(transparency) + + return p + + +def remove_gaussian_particles(name: str, error_if_absent: bool = True) -> None: + """Remove a gaussian particles by name""" + psb.remove_gaussian_particles(name, error_if_absent) + del gaussian_particles_instance_cache[name] + + +def get_gaussian_particles(name: str) -> GaussianParticles: + """Get gaussian particles by name""" + if not has_gaussian_particles(name): + raise ValueError("no gaussian particles with name " + str(name)) + + if name not in gaussian_particles_instance_cache: + # this _should_ always be in sync with the Polyscope C++ structure pool, so if this ever happens something has gone wrong. + raise RuntimeError("internal error: gaussian particles instance cache is missing an entry for " + str(name)) + + return gaussian_particles_instance_cache[name] + + +def has_gaussian_particles(name: str) -> bool: # + """Check if a gaussian particles exists by name""" + return psb.has_gaussian_particles(name) diff --git a/test/scripts/gaussians_demo.py b/test/scripts/gaussians_demo.py new file mode 100644 index 0000000..033e251 --- /dev/null +++ b/test/scripts/gaussians_demo.py @@ -0,0 +1,123 @@ +import torch +import os +import sys +import os.path as path + +# Path to where the bindings live +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src"))) +if os.name == 'nt': # if Windows + # handle default location where VS puts binary + sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "build", "Debug"))) +else: + # normal / unix case + sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "build"))) + + +import polyscope as ps +import polyscope.imgui as psim +from polyscope import imgui as psim + +import sys +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'): + + plydata = PlyData.read(path_ply) + + centers = np.stack((np.asarray(plydata.elements[0]["x"]), + np.asarray(plydata.elements[0]["y"]), + np.asarray(plydata.elements[0]["z"])), axis=1) + opacities = np.asarray(plydata.elements[0]["opacity"]) + + features_dc = np.zeros((centers.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) + + extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] + extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1])) + features_extra = np.zeros((centers.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] + scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1])) + scales = np.zeros((centers.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] + rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1])) + rots = np.zeros((centers.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + with torch.no_grad(): + centers = torch.tensor(centers, dtype=torch.float, device=device,) + features_dc = torch.tensor(features_dc, dtype=torch.float, device=device).transpose(1, 2).contiguous() + # features_rest = torch.tensor(features_extra, dtype=torch.float, device=device).transpose(1, 2).contiguous() + opacity = torch.tensor(opacities, dtype=torch.float, device=device) + scaling = torch.tensor(scales, dtype=torch.float, device=device) + rotation = torch.tensor(rots, dtype=torch.float, device=device) + + + # apply activations + # (assumes file holds pre-activated values) + opacity = torch.sigmoid(opacity) + scaling = torch.exp(scaling) + + return centers, features_dc, opacity, scaling, rotation + +def main(): + + parser = argparse.ArgumentParser() + + # Build arguments + parser.add_argument('--gaussian_ply', type=str, help='path to a .ply file') + + # Parse arguments + args = parser.parse_args() + + + ps.init() + ps.set_ground_plane_mode("shadow_only") + + 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) + ) + + ps.show() + + # ps.remove_gaussian_particles("gaussians") + + +if __name__ == '__main__': + main()