Visualizing gaussian splats

cpp/open3d/t/geometry/PointCloud.cpp

@@ -1369,6 +1369,78 @@ core::Tensor PointCloud::ComputeMetrics(const PointCloud &pcd2,
                                 metrics, params);
 }
 
+bool PointCloud::IsGaussianSplat() const {
+    std::vector<std::string> keys_to_check = {"f_dc", "opacity", "rot", "scale"};
+    for (const auto& key : keys_to_check) {
+        if (point_attr_.find(key) == point_attr_.end()) {
+            return false;
+        }
+    }
+    return true;
+}
+
+int PointCloud::GaussianSplatGetSHOrder() const {
+    if (point_attr_.find("f_rest") == point_attr_.end()) {
+        return 0;
+    } else {
+        const core::Tensor & f_rest = GetPointAttr("f_rest");
+        auto shape = f_rest.GetShape();
+        size_t num_dims = shape.size();
+
+        int dim_of_degree = shape[1];
+        if (num_dims == 2) {
+            if (dim_of_degree % 3 != 0) {
+                utility::LogError("SH coeffs are in shape [N, Nc], but the dimension of Nc cannot be divided by 3");
+            }
+            dim_of_degree = dim_of_degree / 3;
+        }
+
+        std::optional<int> degree = findPositiveIntegerSolution(dim_of_degree);
+        if (degree.has_value()) {
+            if (degree.value() > 2) {
+                utility::LogWarning("The SH coeffs degree higher than 2 is not supported for now.");
+            }
+            return degree.value();
+        } else {
+            utility::LogError("The number of SH coeffs is not a valid value.");
+        }
+        return -1;
+    }
+}
+
+std::optional<int> PointCloud::findPositiveIntegerSolution(int x) const {
+    // Coefficients of the quadratic equation n^2 + 2n - x = 0
+    int a = 1;
+    int b = 2;
+    int c = -x;
+
+    // Calculate the discriminant
+    int discriminant = b * b - 4 * a * c;
+
+    // Check if the discriminant is non-negative and a perfect square
+    if (discriminant < 0) {
+        return std::nullopt;
+    }
+
+    int sqrt_discriminant = static_cast<int>(std::sqrt(discriminant));
+    if (sqrt_discriminant * sqrt_discriminant != discriminant) {
+        return std::nullopt;
+    }
+
+    // Calculate the two possible solutions for n
+    int n1 = (-b + sqrt_discriminant) / (2 * a);
+    int n2 = (-b - sqrt_discriminant) / (2 * a);
+
+    // Check if either solution is a valid positive integer
+    // There is definitely a positive result because x is positive.
+    if (n1 > 0 && n1 * (n1 + 2) == x) {
+        return n1;
+    } else if (n2 > 0 && n2 * (n2 + 2) == x) {
+        return n2;
+    }
+
+    return std::nullopt;
+}
 }  // namespace geometry
 }  // namespace t
 }  // namespace open3d

cpp/open3d/t/geometry/PointCloud.h

@@ -738,9 +738,21 @@ class PointCloud : public Geometry, public DrawableGeometry {
             std::vector<Metric> metrics = {Metric::ChamferDistance},
             MetricParameters params = MetricParameters()) const;
 
+    /// \brief Checks if the point cloud is formatted for Gaussian Splatting.
+    ///
+    /// This function verifies that the point attribute container contains all the required keys
+    /// ("f_dc", "opacity", "rot", and "scale") necessary for Gaussian Splatting. If any of these keys
+    /// are missing, the function returns \c false; otherwise, it returns \c true.
+    /// \return \c true if all required attributes are present, \c false otherwise.
+    bool IsGaussianSplat() const;
+
+    int GaussianSplatGetSHOrder() const;
 protected:
     core::Device device_ = core::Device("CPU:0");
     TensorMap point_attr_;
+
+private:
+    std::optional<int> findPositiveIntegerSolution(int x) const;
 };
 
 }  // namespace geometry

cpp/open3d/visualization/CMakeLists.txt

@@ -80,6 +80,7 @@ if (BUILD_GUI)
         rendering/filament/LineSetBuffers.cpp
         rendering/filament/PointCloudBuffers.cpp
         rendering/filament/TriangleMeshBuffers.cpp
+        rendering/filament/GaussianSplatBuffers.cpp
     )
 
     target_sources(visualization PRIVATE

cpp/open3d/visualization/rendering/MaterialRecord.h

@@ -83,6 +83,9 @@ struct MaterialRecord {
     // Infinite ground plane
     float ground_plane_axis = 0.f;  // 0: XZ; >0: XY; <0: YZ
 
+    // This is only used in gaussian splat.
+    int sh_degree = 2;
+
     // Generic material properties
     std::unordered_map<std::string, Eigen::Vector4f> generic_params;
     std::unordered_map<std::string, geometry::Image> generic_imgs;

cpp/open3d/visualization/rendering/filament/FilamentGeometryBuffersBuilder.cpp

@@ -250,9 +250,14 @@ std::unique_ptr<GeometryBuffersBuilder> GeometryBuffersBuilder::GetBuilder(
     using GT = t::geometry::Geometry::GeometryType;
 
     switch (geometry.GetGeometryType()) {
-        case GT::PointCloud:
-            return std::make_unique<TPointCloudBuffersBuilder>(
-                    static_cast<const t::geometry::PointCloud&>(geometry));
+        case GT::PointCloud: {
+            const t::geometry::PointCloud& pointcloud = static_cast<const t::geometry::PointCloud&>(geometry);
+            if (pointcloud.IsGaussianSplat()) {
+                return std::make_unique<TGaussianSplatBuffersBuilder>(pointcloud);
+            } else {
+                return std::make_unique<TPointCloudBuffersBuilder>(pointcloud);
+            }
+        }
         case GT::TriangleMesh:
             return std::make_unique<TMeshBuffersBuilder>(
                     static_cast<const t::geometry::TriangleMesh&>(geometry));

cpp/open3d/visualization/rendering/filament/FilamentGeometryBuffersBuilder.h

@@ -172,10 +172,40 @@ class TPointCloudBuffersBuilder : public GeometryBuffersBuilder {
     Buffers ConstructBuffers() override;
     filament::Box ComputeAABB() override;
 
-private:
+protected:
     t::geometry::PointCloud geometry_;
 };
 
+class TGaussianSplatBuffersBuilder : public TPointCloudBuffersBuilder {
+public:
+    /// \brief Constructs a TGaussianSplatBuffersBuilder object.
+    ///
+    /// Initializes the Gaussian Splat buffers from the provided \p geometry and ensures that
+    /// all necessary attributes are present and correctly formatted. If the geometry is not a
+    /// Gaussian Splat, a warning is issued. Additionally, attributes like "f_dc", "opacity",
+    /// "rot", "scale", and "f_rest" are checked for their data type, and converted to Float32
+    /// if they are not already in that format.
+    explicit TGaussianSplatBuffersBuilder(const t::geometry::PointCloud& geometry);
+
+    /// \brief Constructs vertex and index buffers for Gaussian Splat rendering.
+    ///
+    /// This function creates and configures GPU buffers to represent a Gaussian Splat
+    /// point cloud. It extracts attributes like positions, colors, rotation, scale,
+    /// and spherical harmonics coefficients from the provided \ref geometry_ and
+    /// organizes them into separate vertex buffer attributes.
+    ///
+    /// The vertex buffer contains the following attributes:
+    /// - POSITION: Vertex positions (FLOAT3)
+    /// - COLOR: DC component and opacity (FLOAT4)
+    /// - TANGENTS: Rotation quaternion (FLOAT4)
+    /// - CUSTOM0: Scale (FLOAT4)
+    /// - CUSTOM1 to CUSTOM6: SH coefficients (FLOAT4)
+    ///
+    /// Each attribute is checked and converted to the expected data type if necessary,
+    /// and missing attributes are initialized with default values.
+    Buffers ConstructBuffers() override;
+};
+
 class TLineSetBuffersBuilder : public GeometryBuffersBuilder {
 public:
     explicit TLineSetBuffersBuilder(const t::geometry::LineSet& geometry);

cpp/open3d/visualization/rendering/filament/FilamentResourceManager.cpp

@@ -318,6 +318,8 @@ TextureSettings GetSettingsFromImage(const t::geometry::Image& image,
 
 const MaterialHandle FilamentResourceManager::kDefaultLit =
         MaterialHandle::Next();
+const MaterialHandle FilamentResourceManager::kGaussianSplatShader =
+        MaterialHandle::Next();
 const MaterialHandle FilamentResourceManager::kDefaultLitWithTransparency =
         MaterialHandle::Next();
 const MaterialHandle FilamentResourceManager::kDefaultLitSSR =
@@ -1017,6 +1019,25 @@ void FilamentResourceManager::LoadDefaults() {
     lit_mat->setDefaultParameter("anisotropyMap", texture, default_sampler);
     materials_[kDefaultLit] = BoxResource(lit_mat, engine_);
 
+    const auto gaussian_path = resource_root + "/gaussianSplat.filamat";
+    auto gaussian_mat = LoadMaterialFromFile(gaussian_path, engine_);
+    gaussian_mat->setDefaultParameter("baseColor", filament::RgbType::sRGB,
+                                    default_color);
+    gaussian_mat->setDefaultParameter("baseRoughness", 0.7f);
+    gaussian_mat->setDefaultParameter("reflectance", 0.5f);
+    gaussian_mat->setDefaultParameter("baseMetallic", 0.f);
+    gaussian_mat->setDefaultParameter("clearCoat", 0.f);
+    gaussian_mat->setDefaultParameter("clearCoatRoughness", 0.f);
+    gaussian_mat->setDefaultParameter("anisotropy", 0.f);
+    gaussian_mat->setDefaultParameter("pointSize", 3.f);
+    gaussian_mat->setDefaultParameter("albedo", texture, default_sampler);
+    gaussian_mat->setDefaultParameter("ao_rough_metalMap", texture, default_sampler);
+    gaussian_mat->setDefaultParameter("normalMap", normal_map, default_sampler);
+    gaussian_mat->setDefaultParameter("reflectanceMap", texture, default_sampler);
+
+    gaussian_mat->setDefaultParameter("anisotropyMap", texture, default_sampler);
+    materials_[kGaussianSplatShader] = BoxResource(gaussian_mat, engine_);
+
     const auto lit_trans_path =
             resource_root + "/defaultLitTransparency.filamat";
     auto lit_trans_mat = LoadMaterialFromFile(lit_trans_path, engine_);

cpp/open3d/visualization/rendering/filament/FilamentResourceManager.h

@@ -50,6 +50,7 @@ namespace rendering {
 class FilamentResourceManager {
 public:
     static const MaterialHandle kDefaultLit;
+    static const MaterialHandle kGaussianSplatShader;
     static const MaterialHandle kDefaultLitWithTransparency;
     static const MaterialHandle kDefaultLitSSR;
     static const MaterialHandle kDefaultUnlit;

cpp/open3d/visualization/rendering/filament/FilamentScene.cpp

@@ -102,7 +102,8 @@ std::unordered_map<std::string, MaterialHandle> shader_mappings = {
          ResourceManager::kDefaultUnlitPolygonOffsetShader},
         {"unlitBackground", ResourceManager::kDefaultUnlitBackgroundShader},
         {"infiniteGroundPlane", ResourceManager::kInfinitePlaneShader},
-        {"unlitLine", ResourceManager::kDefaultLineShader}};
+        {"unlitLine", ResourceManager::kDefaultLineShader},
+        {"gaussianSplat", ResourceManager::kGaussianSplatShader}};
 
 MaterialHandle kColorOnlyMesh = ResourceManager::kDefaultUnlit;
 MaterialHandle kPlainMesh = ResourceManager::kDefaultLit;
@@ -846,6 +847,32 @@ void FilamentScene::UpdateDefaultLit(GeometryMaterialInstance& geom_mi) {
             .Finish();
 }
 
+void FilamentScene::UpdateGaussianSplat(GeometryMaterialInstance& geom_mi) {
+    auto& material = geom_mi.properties;
+    auto& maps = geom_mi.maps;
+
+    renderer_.ModifyMaterial(geom_mi.mat_instance)
+            .SetColor("baseColor", material.base_color, false)
+            .SetParameter("pointSize", material.point_size)
+            .SetParameter("baseRoughness", material.base_roughness)
+            .SetParameter("baseMetallic", material.base_metallic)
+            .SetParameter("reflectance", material.base_reflectance)
+            .SetParameter("clearCoat", material.base_clearcoat)
+            .SetParameter("clearCoatRoughness",
+                          material.base_clearcoat_roughness)
+            .SetParameter("anisotropy", material.base_anisotropy)
+            .SetParameter("shDegree", material.sh_degree)
+            .SetTexture("albedo", maps.albedo_map,
+                        rendering::TextureSamplerParameters::Pretty())
+            .SetTexture("normalMap", maps.normal_map,
+                        rendering::TextureSamplerParameters::Pretty())
+            .SetTexture("ao_rough_metalMap", maps.ao_rough_metal_map,
+                        rendering::TextureSamplerParameters::Pretty())
+            .SetTexture("reflectanceMap", maps.reflectance_map,
+                        rendering::TextureSamplerParameters::Pretty())
+            .Finish();
+}
+
 void FilamentScene::UpdateDefaultLitSSR(GeometryMaterialInstance& geom_mi) {
     auto& material = geom_mi.properties;
     auto& maps = geom_mi.maps;
@@ -1131,6 +1158,8 @@ void FilamentScene::UpdateMaterialProperties(RenderableGeometry& geom) {
         UpdateLineShader(geom.mat);
     } else if (props.shader == "unlitPolygonOffset") {
         UpdateUnlitPolygonOffsetShader(geom.mat);
+    } else if (props.shader == "gaussianSplat") {
+        UpdateGaussianSplat(geom.mat);
     } else {
         utility::LogWarning("'{}' is not a valid shader", props.shader);
     }

cpp/open3d/visualization/rendering/filament/FilamentScene.h

@@ -305,6 +305,7 @@ class FilamentScene : public Scene {
                                   bool shader_only = false);
     void UpdateMaterialProperties(RenderableGeometry& geom);
     void UpdateDefaultLit(GeometryMaterialInstance& geom_mi);
+    void UpdateGaussianSplat(GeometryMaterialInstance& geom_mi);
     void UpdateDefaultLitSSR(GeometryMaterialInstance& geom_mi);
     void UpdateDefaultUnlit(GeometryMaterialInstance& geom_mi);
     void UpdateNormalShader(GeometryMaterialInstance& geom_mi);