Adding support for .spz files

CMakeLists.txt

@@ -57,7 +57,15 @@ FetchContent_Declare(nanoflann
 FetchContent_Declare(cxxopts
     URL https://github.com/jarro2783/cxxopts/archive/refs/tags/v3.2.0.zip
 )
-FetchContent_MakeAvailable(nlohmann_json nanoflann cxxopts)
+FetchContent_Declare(spz
+    URL https://github.com/nianticlabs/spz/archive/refs/tags/v3.0.0.zip
+)
+set(SPZ_BUILD_PYTHON_BINDINGS OFF)
+set(SPZ_BUILD_TOOLS OFF)
+set(SPZ_BUILD_EXTENSIONS OFF)
+set(SPZ_BUILD_WASM OFF)
+set(SPZ_SHOULD_INSTALL OFF)
+FetchContent_MakeAvailable(nlohmann_json nanoflann cxxopts spz)
 if((GPU_RUNTIME STREQUAL "CUDA") OR (GPU_RUNTIME STREQUAL "HIP"))
     FetchContent_Declare(glm
         URL https://github.com/g-truc/glm/archive/refs/tags/1.0.1.zip
@@ -267,6 +275,7 @@ target_link_libraries(opensplat PRIVATE
     nlohmann_json::nlohmann_json
     cxxopts::cxxopts
     nanoflann::nanoflann
+    spz::spz
 )
 if (NOT WIN32)
     target_link_libraries(opensplat PUBLIC pthread)

README.md

@@ -5,7 +5,7 @@ A free and open source implementation of 3D [gaussian splatting](https://www.you
 <img src="https://github.com/pierotofy/OpenSplat/assets/1951843/c9327c7c-31ad-402d-a5a5-04f7602ca5f5" width="49%" />
 <img src="https://github.com/pierotofy/OpenSplat/assets/1951843/eba4ae75-2c88-4c9e-a66b-608b574d085f" width="49%" />
 
-OpenSplat takes camera poses + sparse points in [COLMAP](https://colmap.github.io/), [OpenSfM](https://github.com/mapillary/OpenSfM), [ODM](https://github.com/OpenDroneMap/ODM),  [OpenMVG](https://github.com/OpenMVG/OpenMVG) or [nerfstudio](https://docs.nerf.studio/quickstart/custom_dataset.html) project format and computes a [scene file](https://drive.google.com/file/d/12lmvVWpFlFPL6nxl2e2d-4u4a31RCSKT/view?usp=sharing) (.ply or .splat) that can be later imported for [viewing](https://antimatter15.com/splat/?url=https://splat.uav4geo.com/banana.splat), editing and rendering in other [software](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#open-source-implementations).
+OpenSplat takes camera poses + sparse points in [COLMAP](https://colmap.github.io/), [OpenSfM](https://github.com/mapillary/OpenSfM), [ODM](https://github.com/OpenDroneMap/ODM),  [OpenMVG](https://github.com/OpenMVG/OpenMVG) or [nerfstudio](https://docs.nerf.studio/quickstart/custom_dataset.html) project format and computes a [scene file](https://drive.google.com/file/d/12lmvVWpFlFPL6nxl2e2d-4u4a31RCSKT/view?usp=sharing) (.ply, .splat, or .spz) that can be later imported for [viewing](https://antimatter15.com/splat/?url=https://splat.uav4geo.com/banana.splat), editing and rendering in other [software](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#open-source-implementations).
 
 Graphics card recommended, but not required! OpenSplat runs the fastest on NVIDIA, AMD and Apple (Metal) GPUs, but can also run entirely on the CPU (~100x slower).
 

model.cpp

@@ -1,12 +1,14 @@
 #include <filesystem>
 #include "model.hpp"
 #include "constants.hpp"
+#include "splat-types.h"
 #include "tile_bounds.hpp"
 #include "project_gaussians.hpp"
 #include "rasterize_gaussians.hpp"
 #include "tensor_math.hpp"
 #include "gsplat.hpp"
 #include "utils.hpp"
+#include <load-spz.h>
 
 #ifdef USE_MPS
 #include <torch/mps.h>
@@ -55,6 +57,11 @@ torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){
     return torch::abs(gt - rendered).mean();
 }
 
+template<typename T>
+std::vector<T> tensor_to_vector(const torch::Tensor t){
+    return std::vector<T>(t.data_ptr<float>(), t.data_ptr<float>() + t.numel());
+}
+
 void Model::setupOptimizers(){
     releaseOptimizers();
 
@@ -494,12 +501,24 @@ void Model::afterTrain(int step){
 }
 
 void Model::save(const std::string &filename, int step){
-    if (fs::path(filename).extension().string() == ".splat"){
+    std::string extension = fs::path(filename).extension().string(); 
+    if (extension == ".splat"){
         saveSplat(filename);
-    }else{
+        std::cout << "Wrote " << filename << std::endl;
+    }
+    else if (extension == ".ply") {
         savePly(filename, step);
+        std::cout << "Wrote " << filename << std::endl;
+    }
+    else {
+        bool success = saveSpz(filename);
+        if (success) {
+            std::cout << "Wrote " << filename << std::endl;
+        }
+        else {
+            std::cerr << "Failed to write " << filename << ", aborting save." << std::endl;
+        }
     }
-    std::cout << "Wrote " << filename << std::endl;
 }
 
 void Model::savePly(const std::string &filename, int step){
@@ -597,6 +616,41 @@ void Model::saveSplat(const std::string &filename){
     o.close();
 }
 
+bool Model::saveSpz(const std::string &filename){
+    auto numPoints = means.size(0);
+
+
+    torch::Tensor meansCpu = keepCrs ? (means.cpu() / scale) + translation : means.cpu();
+    torch::Tensor scalesCpu = keepCrs ? (scales.cpu() / scale) : scales.cpu();
+
+    torch::Tensor meansFlat = meansCpu.flatten();
+    torch::Tensor scalesFlat = scalesCpu.flatten();
+    torch::Tensor colorsFlat = featuresDc.cpu().flatten(); // raw DC coefficients
+    torch::Tensor opacFlat = opacities.flatten().cpu();
+    torch::Tensor quatsFlat = quats.flatten().cpu();
+    torch::Tensor shRestFlat = featuresRest.cpu().transpose(1, 2).flatten();
+
+    spz::GaussianCloud gaussians;
+    gaussians.numPoints = meansCpu.size(0);
+    gaussians.shDegree = shDegree;
+    gaussians.antialiased = false;
+    gaussians.positions = tensor_to_vector<float>(meansFlat);
+    gaussians.scales = tensor_to_vector<float>(scalesFlat);
+    gaussians.rotations = tensor_to_vector<float>(quatsFlat);
+    gaussians.alphas = tensor_to_vector<float>(opacFlat);
+    gaussians.colors = tensor_to_vector<float>(colorsFlat);
+    gaussians.sh = tensor_to_vector<float>(shRestFlat);
+
+    auto options = spz::PackOptions{
+        .version = 3,       // V4 available but not handled by many viewers
+        .from = spz::CoordinateSystem::RUB,
+        .sh1Bits = 6,
+        .shRestBits = 5
+    };
+    bool success = spz::saveSpz(gaussians, options, filename);
+    return success;
+}
+
 void Model::saveDebugPly(const std::string &filename, int step){
     // A standard PLY
     std::ofstream o(filename, std::ios::binary);

model.hpp

@@ -72,6 +72,7 @@ struct Model{
   void save(const std::string &filename, int step);
   void savePly(const std::string &filename, int step);
   void saveSplat(const std::string &filename);
+  bool saveSpz(const std::string &filename);
   void saveDebugPly(const std::string &filename, int step);
   int loadPly(const std::string &filename);
   torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight);