fixed the cmake

Author: eokeeffe

CMakeLists.txt

@@ -5,6 +5,7 @@ set(OPENSPLAT_BUILD_SIMPLE_TRAINER OFF CACHE BOOL "Build simple trainer applicat
 set(GPU_RUNTIME "CUDA" CACHE STRING "HIP or CUDA or MPS")
 set(OPENCV_DIR "OPENCV_DIR-NOTFOUND" CACHE PATH "Path to the OPENCV installation directory")
 set(OPENSPLAT_MAX_CUDA_COMPATIBILITY OFF CACHE BOOL "Build for maximum CUDA device compatibility")
+set(OPENSPLAT_BUILD_VISUALIZER OFF CACHE BOOL "Build visualizer application")
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
 
 # Read version
@@ -163,6 +164,9 @@ endif()
 
 find_package(Torch REQUIRED)
 find_package(OpenCV HINTS "${OPENCV_DIR}" REQUIRED)
+if (OPENSPLAT_BUILD_VISUALIZER)
+    find_package(Pangolin REQUIRED)
+endif()
 
 if (NOT WIN32 AND NOT APPLE)
     set(CMAKE_CUDA_COMPILER "${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc")
@@ -216,14 +220,33 @@ endif()
 add_library(gsplat_cpu rasterizer/gsplat-cpu/gsplat_cpu.cpp)
 target_include_directories(gsplat_cpu PRIVATE ${TORCH_INCLUDE_DIRS})
 
-add_executable(opensplat opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp tensor_math.cpp)
+set(OPENSPLAT_SRC_FILES opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp
+kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp
+rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp
+tensor_math.cpp)
+
+if (OPENSPLAT_BUILD_VISUALIZER)
+    if (Pangolin_FOUND)
+        message(STATUS "Found Pangolin. Building visualizer (beta)")
+        list(APPEND OPENSPLAT_SRC_FILES visualizer.cpp)
+        add_definitions(-DUSE_VISUALIZATION)
+    else()
+        message(FATAL "Pangolin not found. Cannot build visualizer (beta)")
+    endif()
+endif()
+
+add_executable(opensplat ${OPENSPLAT_SRC_FILES})
+
 install(TARGETS opensplat DESTINATION bin)
 set_property(TARGET opensplat PROPERTY CXX_STANDARD 17)
 target_include_directories(opensplat PRIVATE
     ${PROJECT_SOURCE_DIR}/rasterizer
     ${GPU_INCLUDE_DIRS}
 )
 target_link_libraries(opensplat PUBLIC ${STDPPFS_LIBRARY} ${GPU_LIBRARIES} ${GSPLAT_LIBS} ${TORCH_LIBRARIES} ${OpenCV_LIBS})
+if (Pangolin_FOUND)
+    target_link_libraries(opensplat PUBLIC ${Pangolin_LIBRARIES})
+endif()
 target_link_libraries(opensplat PRIVATE
     nlohmann_json::nlohmann_json
     cxxopts::cxxopts

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) 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 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).
 
 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).
 
@@ -140,6 +140,8 @@ If building CPU-only, remove `-DGPU_RUNTIME=MPS`.
 
 :warning: You will probably get a *libc10.dylib can’t be opened because Apple cannot check it for malicious software* error on first run. Open **System Settings** and go to **Privacy & Security** and find the **Allow** button. You might need to repeat this several times until all torch libraries are loaded.
 
+:warning: If you get a *Library not loaded: @rpath/libomp.dylib* error, try running `brew link libomp --force` before running OpenSplat.
+
 ## Docker Build
 
 ### CUDA
@@ -231,6 +233,14 @@ To generate compressed splats (.splat files), use the `-o` option:
 ./opensplat /path/to/banana -o banana.splat
 ```
 
+### Resume
+
+You can resume training of a .PLY file by using the `--resume` option:
+
+```bash
+./opensplat /path/to/banana --resume ./splat.ply
+```
+
 ### AMD GPU Notes
 
 To train a model with AMD GPU using docker container, you can use the following command as a reference:

input_data.cpp

@@ -44,22 +44,25 @@ void Camera::loadImage(float downscaleFactor){
     if (image.numel()) std::runtime_error("loadImage already called");
     std::cout << "Loading " << filePath << std::endl;
 
-    float scaleFactor = 1.0f / downscaleFactor;
     cv::Mat cImg = imreadRGB(filePath);
 
     float rescaleF = 1.0f;
     // If camera intrinsics don't match the image dimensions 
     if (cImg.rows != height || cImg.cols != width){
         rescaleF = static_cast<float>(cImg.rows) / static_cast<float>(height);
     }
-    fx *= scaleFactor * rescaleF;
-    fy *= scaleFactor * rescaleF;
-    cx *= scaleFactor * rescaleF;
-    cy *= scaleFactor * rescaleF;
+    fx *= rescaleF;
+    fy *= rescaleF;
+    cx *= rescaleF;
+    cy *= rescaleF;
 
     if (downscaleFactor > 1.0f){
-        float f = 1.0f / downscaleFactor;
-        cv::resize(cImg, cImg, cv::Size(), f, f, cv::INTER_AREA);
+        float scaleFactor = 1.0f / downscaleFactor;
+        cv::resize(cImg, cImg, cv::Size(), scaleFactor, scaleFactor, cv::INTER_AREA);
+        fx *= scaleFactor;
+        fy *= scaleFactor;
+        cx *= scaleFactor;
+        cy *= scaleFactor;
     }
 
     K = getIntrinsicsMatrix();

model.cpp

@@ -6,6 +6,7 @@
 #include "rasterize_gaussians.hpp"
 #include "tensor_math.hpp"
 #include "gsplat.hpp"
+#include "utils.hpp"
 
 #ifdef USE_HIP
 #include <c10/hip/HIPCachingAllocator.h>
@@ -50,6 +51,30 @@ torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){
     return torch::abs(gt - rendered).mean();
 }
 
+void Model::setupOptimizers(){
+    releaseOptimizers();
+
+    meansOpt = new torch::optim::Adam({means}, torch::optim::AdamOptions(0.00016));
+    scalesOpt = new torch::optim::Adam({scales}, torch::optim::AdamOptions(0.005));
+    quatsOpt = new torch::optim::Adam({quats}, torch::optim::AdamOptions(0.001));
+    featuresDcOpt = new torch::optim::Adam({featuresDc}, torch::optim::AdamOptions(0.0025));
+    featuresRestOpt = new torch::optim::Adam({featuresRest}, torch::optim::AdamOptions(0.000125));
+    opacitiesOpt = new torch::optim::Adam({opacities}, torch::optim::AdamOptions(0.05));
+
+    meansOptScheduler = new OptimScheduler(meansOpt, 0.0000016f, maxSteps);
+}
+
+void Model::releaseOptimizers(){
+    RELEASE_SAFELY(meansOpt);
+    RELEASE_SAFELY(scalesOpt);
+    RELEASE_SAFELY(quatsOpt);
+    RELEASE_SAFELY(featuresDcOpt);
+    RELEASE_SAFELY(featuresRestOpt);
+    RELEASE_SAFELY(opacitiesOpt);
+
+    RELEASE_SAFELY(meansOptScheduler);
+}
+
 
 torch::Tensor Model::forward(Camera& cam, int step){
 
@@ -139,13 +164,11 @@ torch::Tensor Model::forward(Camera& cam, int step){
         #endif
     }
 
+    xys.retain_grad();
 
     if (radii.sum().item<float>() == 0.0f)
         return backgroundColor.repeat({height, width, 1});
 
-    // TODO: is this needed?
-    xys.retain_grad();
-
     torch::Tensor viewDirs = means.detach() - T.transpose(0, 1).to(device);
     viewDirs = viewDirs / viewDirs.norm(2, {-1}, true);
     int degreesToUse = (std::min<int>)(step / shDegreeInterval, shDegree);
@@ -281,6 +304,9 @@ void Model::removeFromOptimizer(torch::optim::Adam *optimizer, const torch::Tens
 void Model::afterTrain(int step){
     torch::NoGradGuard noGrad;
 
+    // When radii.sum() == 0
+    if (!xys.grad().defined()) return;
+
     if (step < stopSplitAt){
         torch::Tensor visibleMask = (radii > 0).flatten();
 
@@ -460,22 +486,22 @@ void Model::afterTrain(int step){
     }
 }
 
-void Model::save(const std::string &filename){
+void Model::save(const std::string &filename, int step){
     if (fs::path(filename).extension().string() == ".splat"){
         saveSplat(filename);
     }else{
-        savePly(filename);
+        savePly(filename, step);
     }
     std::cout << "Wrote " << filename << std::endl;
 }
 
-void Model::savePly(const std::string &filename){
+void Model::savePly(const std::string &filename, int step){
     std::ofstream o(filename, std::ios::binary);
     int numPoints = means.size(0);
 
     o << "ply" << std::endl;
     o << "format binary_little_endian 1.0" << std::endl;
-    o << "comment Generated by opensplat" << std::endl;
+    o << "comment Generated by opensplat at iteration " << step << std::endl;
     o << "element vertex " << numPoints << std::endl;
     o << "property float x" << std::endl;
     o << "property float y" << std::endl;
@@ -564,14 +590,14 @@ void Model::saveSplat(const std::string &filename){
     o.close();
 }
 
-void Model::saveDebugPly(const std::string &filename){
+void Model::saveDebugPly(const std::string &filename, int step){
     // A standard PLY
     std::ofstream o(filename, std::ios::binary);
     int numPoints = means.size(0);
 
     o << "ply" << std::endl;
     o << "format binary_little_endian 1.0" << std::endl;
-    o << "comment Generated by opensplat" << std::endl;
+    o << "comment Generated by opensplat at iteration " << step << std::endl;
     o << "element vertex " << numPoints << std::endl;
     o << "property float x" << std::endl;
     o << "property float y" << std::endl;
@@ -593,6 +619,157 @@ void Model::saveDebugPly(const std::string &filename){
     std::cout << "Wrote " << filename << std::endl;
 }
 
+int Model::loadPly(const std::string &filename){
+    std::ifstream f(filename, std::ios::binary);
+    if (!f.is_open()) throw std::runtime_error("Invalid PLY file");
+
+    // Ensure we have a valid ply file
+    std::string line;
+    int numPoints;
+    int step;
+    size_t bytesRead = 0;
+
+    std::getline(f, line);
+    bytesRead += f.gcount();
+
+    if (line == "ply"){
+        std::getline(f, line);
+        bytesRead += f.gcount();
+        if (line == "format binary_little_endian 1.0"){
+            std::getline(f, line);
+            bytesRead += f.gcount();
+            const std::string pattern = "comment Generated by opensplat at iteration ";
+
+            if (line.rfind(pattern, 0) == 0){
+                step = std::stoi(line.substr(pattern.length()));
+                if (step >= 0){
+                    std::getline(f, line);
+                    bytesRead += f.gcount();
+                    const std::string pattern = "element vertex ";
+
+                    if (line.rfind(pattern, 0) == 0){
+                        const int numPoints = std::stoi(line.substr(pattern.length()));
+                        
+                        const char *requiredProps[] = {
+                            "property float x",
+                            "property float y",
+                            "property float z",
+                            "property float nx",
+                            "property float ny",
+                            "property float nz",
+                            "property float f_dc_"
+                            "property float f_rest_",
+                            "property float opacity",
+                            "property float scale_0",
+                            "property float scale_1",
+                            "property float scale_2",
+                            "property float rot_0",
+                            "property float rot_1",
+                            "property float rot_2",
+                            "property float rot_3",
+                            "end_header"
+                        };
+
+                        for (int i = 0; i < 6; i++){
+                            std::getline(f, line);
+                            bytesRead += f.gcount();
+                            if (line != requiredProps[i]){
+                                throw std::runtime_error(std::string("PLY file's header does not contain required property: ") + requiredProps[i]);
+                            }
+                        }
+                        std::getline(f, line);
+                        bytesRead += f.gcount();
+
+                        auto countPrefixes = [&f, &line](const char *prefix){
+                                int n = 0;
+                                while(true){
+                                    if (line.rfind(prefix, 0) == 0){
+                                        ++n;
+                                        std::getline(f, line);
+                                    } else {
+                                        break;
+                                    }
+                                }
+                                return n;
+                        };
+                        int featuresDcSize = countPrefixes("property float f_dc_");
+                        int featuresRestSize = countPrefixes("property float f_rest_");
+                        
+                        bool foundEnd = false;
+                        for (int i = 8; i < std::size(requiredProps); i++){
+                            std::getline(f, line);
+                            bytesRead += f.gcount();
+
+                            if (line != requiredProps[i]){
+                                throw std::runtime_error(std::string("PLY file's header does not contain required property: ") + requiredProps[i]);
+                            }
+
+                            if (line == "end_header"){
+                                foundEnd = true;
+                                break;
+                            }
+                        }
+
+                        if (!foundEnd){
+                            throw std::runtime_error("PLY file header does not contain header end");
+                        } 
+
+                        const size_t bytesPerPoint = sizeof(float) * (14 + featuresDcSize + featuresRestSize);
+                        const size_t remainingFileSize = fs::file_size(filename) - bytesRead;
+                        if (remainingFileSize != bytesPerPoint * numPoints){
+                            std::cout << "Loading PLY..." << std::endl;
+                            
+                            float zeros[3];
+
+                            torch::Tensor meansCpu = torch::zeros({numPoints, 3}, torch::TensorOptions().dtype(torch::kFloat32));
+                            torch::Tensor featuresDcCpu = torch::zeros({numPoints, featuresDcSize}, torch::TensorOptions().dtype(torch::kFloat32));
+                            torch::Tensor featuresRestCpu = torch::zeros({numPoints, featuresRestSize}, torch::TensorOptions().dtype(torch::kFloat32));
+                            torch::Tensor opacitiesCpu = torch::zeros({numPoints, 1}, torch::TensorOptions().dtype(torch::kFloat32));
+                            torch::Tensor scalesCpu = torch::zeros({numPoints, 3}, torch::TensorOptions().dtype(torch::kFloat32));
+                            torch::Tensor quatsCpu = torch::zeros({numPoints, 4}, torch::TensorOptions().dtype(torch::kFloat32));
+
+                            for (size_t i = 0; i < numPoints; i++){
+                                f.read(reinterpret_cast<char *>(meansCpu[i].data_ptr()), sizeof(float) * 3);
+                                f.read(reinterpret_cast<char *>(&zeros[0]), sizeof(float) * 3);
+                                f.read(reinterpret_cast<char *>(featuresDcCpu[i].data_ptr()), sizeof(float) * featuresDcSize);
+                                f.read(reinterpret_cast<char *>(featuresRestCpu[i].data_ptr()), sizeof(float) * featuresRestSize);
+                                f.read(reinterpret_cast<char *>(opacitiesCpu[i].data_ptr()), sizeof(float) * 1);
+                                f.read(reinterpret_cast<char *>(scalesCpu[i].data_ptr()), sizeof(float) * 3);
+                                f.read(reinterpret_cast<char *>(quatsCpu[i].data_ptr()), sizeof(float) * 4);
+                            }
+                            if (keepCrs){
+                                meansCpu = (meansCpu - translation) * scale;
+                                scalesCpu = torch::log(scale * torch::exp(scalesCpu));
+                            }
+                            
+                            means = meansCpu.to(device).requires_grad_();
+                            featuresDc = featuresDcCpu.to(device).requires_grad_();
+                            featuresRest = featuresRestCpu.reshape({numPoints, 3, featuresRestSize/3}).transpose(2, 1).to(device).requires_grad_();
+                            opacities = opacitiesCpu.to(device).requires_grad_();
+                            scales = scalesCpu.to(device).requires_grad_();
+                            quats = quatsCpu.to(device).requires_grad_();
+                            
+                            std::cerr << "Loaded " << means.size(0) << " gaussians" << std::endl;
+                            
+                            setupOptimizers();
+                            
+                            f.close();
+                            return step;
+                        } else {
+                            throw std::runtime_error("PLY file's data section is wrong size");
+                        }
+                    }
+                } else {
+                    throw std::runtime_error("PLY file failed sanity check: iteration count should not begin at 0");
+                }
+            } else if (line.rfind("comment Generated by opensplat")){
+                throw std::runtime_error("PLY file does not contain iteration count metadata. You can edit the file to add this metadata manually, by changing \"comment Generated by opensplat\" to \"comment Generated by opensplat at iteration 12345\", changing 12345 to the desired value.");
+            }
+        }
+    }
+    throw std::runtime_error("Invalid PLY file");
+}
+
 torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight){
     torch::Tensor ssimLoss = 1.0f - ssim.eval(rgb, gt);
     torch::Tensor l1Loss = l1(rgb, gt);