From 4f358c99a7544f76128d785da082cbfbcc83df08 Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Mon, 25 Mar 2024 17:55:13 -0400
Subject: [PATCH] Remove stop-split-at (auto compute), stop culling when
 stopping splits

---
 model.cpp     | 2 +-
 model.hpp     | 4 ++--
 opensplat.cpp | 4 +---
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/model.cpp b/model.cpp
index 81a7ccf..f5461a7 100644
--- a/model.cpp
+++ b/model.cpp
@@ -387,7 +387,7 @@ void Model::afterTrain(int step){
             std::cout << "Added " << (means.size(0) - numPointsBefore) << " gaussians, new count " << means.size(0) << std::endl;
         }
 
-        if (doDensification || step >= stopSplitAt){
+        if (doDensification){
             // Cull
             int numPointsBefore = means.size(0);
 
diff --git a/model.hpp b/model.hpp
index e78b1e4..74a22ad 100644
--- a/model.hpp
+++ b/model.hpp
@@ -22,12 +22,12 @@ torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt);
 struct Model{
   Model(const InputData &inputData, int numCameras,
         int numDownscales, int resolutionSchedule, int shDegree, int shDegreeInterval, 
-        int refineEvery, int warmupLength, int resetAlphaEvery, int stopSplitAt, float densifyGradThresh, float densifySizeThresh, int stopScreenSizeAt, float splitScreenSize,
+        int refineEvery, int warmupLength, int resetAlphaEvery, float densifyGradThresh, float densifySizeThresh, int stopScreenSizeAt, float splitScreenSize,
         int maxSteps,
         const torch::Device &device) :
     numCameras(numCameras),
     numDownscales(numDownscales), resolutionSchedule(resolutionSchedule), shDegree(shDegree), shDegreeInterval(shDegreeInterval), 
-    refineEvery(refineEvery), warmupLength(warmupLength), resetAlphaEvery(resetAlphaEvery), stopSplitAt(stopSplitAt), densifyGradThresh(densifyGradThresh), densifySizeThresh(densifySizeThresh), stopScreenSizeAt(stopScreenSizeAt), splitScreenSize(splitScreenSize),
+    refineEvery(refineEvery), warmupLength(warmupLength), resetAlphaEvery(resetAlphaEvery), stopSplitAt(maxSteps / 2), densifyGradThresh(densifyGradThresh), densifySizeThresh(densifySizeThresh), stopScreenSizeAt(stopScreenSizeAt), splitScreenSize(splitScreenSize),
     maxSteps(maxSteps),
     device(device), ssim(11, 3){
 
diff --git a/opensplat.cpp b/opensplat.cpp
index e0fc60c..9cbf4b0 100644
--- a/opensplat.cpp
+++ b/opensplat.cpp
@@ -30,7 +30,6 @@ int main(int argc, char *argv[]){
         ("refine-every", "Split/duplicate/prune gaussians every these many steps", cxxopts::value<int>()->default_value("100"))
         ("warmup-length", "Split/duplicate/prune gaussians only after these many steps", cxxopts::value<int>()->default_value("500"))
         ("reset-alpha-every", "Reset the opacity values of gaussians after these many refinements (not steps)", cxxopts::value<int>()->default_value("30"))
-        ("stop-split-at", "Stop splitting/duplicating gaussians after these many steps", cxxopts::value<int>()->default_value("15000"))
         ("densify-grad-thresh", "Threshold of the positional gradient norm (magnitude of the loss function) which when exceeded leads to a gaussian split/duplication", cxxopts::value<float>()->default_value("0.0002"))
         ("densify-size-thresh", "Gaussians' scales below this threshold are duplicated, otherwise split", cxxopts::value<float>()->default_value("0.01"))
         ("stop-screen-size-at", "Stop splitting gaussians that are larger than [split-screen-size] after these many steps", cxxopts::value<int>()->default_value("4000"))
@@ -73,7 +72,6 @@ int main(int argc, char *argv[]){
     const int refineEvery = result["refine-every"].as<int>();
     const int warmupLength = result["warmup-length"].as<int>();
     const int resetAlphaEvery = result["reset-alpha-every"].as<int>();
-    const int stopSplitAt = result["stop-split-at"].as<int>();
     const float densifyGradThresh = result["densify-grad-thresh"].as<float>();
     const float densifySizeThresh = result["densify-size-thresh"].as<float>();
     const int stopScreenSizeAt = result["stop-screen-size-at"].as<int>();
@@ -104,7 +102,7 @@ int main(int argc, char *argv[]){
         Model model(inputData,
                     cams.size(),
                     numDownscales, resolutionSchedule, shDegree, shDegreeInterval, 
-                    refineEvery, warmupLength, resetAlphaEvery, stopSplitAt, densifyGradThresh, densifySizeThresh, stopScreenSizeAt, splitScreenSize,
+                    refineEvery, warmupLength, resetAlphaEvery, densifyGradThresh, densifySizeThresh, stopScreenSizeAt, splitScreenSize,
                     numIters,
                     device);