From 0327c00835b16cea1a503f0517ab5f89f2134cdc Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Sat, 17 Feb 2024 14:27:13 -0500
Subject: [PATCH] More verbose, update readme

---
 README.md     |  3 +++
 model.cpp     |  8 +++++---
 opensplat.cpp | 10 ++++++++--
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 47d6bf1..d036a4f 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,9 @@
 
 A free and open source implementation of 3D gaussian splatting, written in C++. It's based on [splatfacto](https://docs.nerf.studio/nerfology/methods/splat.html) and focuses on being portable, lean and fast.
 
+![OpenSplat](https://github.com/pierotofy/OpenSplat/assets/1951843/3461e0e4-e134-4d6a-8a56-d89d00258e41)
+
+
 OpenSplat takes camera poses + sparse points and computes a scene file (.ply) that can be later imported for viewing, editing and rendering in other [software](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#open-source-implementations).
 
 ## Build
diff --git a/model.cpp b/model.cpp
index 8be917d..b2e3a22 100644
--- a/model.cpp
+++ b/model.cpp
@@ -255,6 +255,7 @@ void Model::afterTrain(int step){
         const float cullAlphaThresh = 0.1f;
 
         if (doDensification){
+            int numPointsBefore = means.size(0);
             torch::Tensor avgGradNorm = (xysGradNorm / visCounts) * 0.5f * static_cast<float>( (std::max)(lastWidth, lastHeight) );
             torch::Tensor highGrads = (avgGradNorm > densifyGradThresh).squeeze();
 
@@ -288,7 +289,6 @@ void Model::afterTrain(int step){
             // Duplicate gaussians that are too small
             torch::Tensor dups = (std::get<0>(scales.exp().max(-1)) <= densifySizeThresh).squeeze();
             dups &= highGrads;
-            int nDups = dups.sum().item<int>();
             torch::Tensor dupMeans = means.index({dups});
             torch::Tensor dupFeaturesDc = featuresDc.index({dups});
             torch::Tensor dupFeaturesRest = featuresRest.index({dups});
@@ -330,6 +330,8 @@ void Model::afterTrain(int step){
                 splits,
                 torch::full({nSplitSamples * splits.sum().item<int>() + dups.sum().item<int>()}, false, torch::TensorOptions().dtype(torch::kBool).device(device))
             }, 0);
+
+            std::cout << "Added " << (means.size(0) - numPointsBefore) << " gaussians, new count " << means.size(0) << std::endl;
         }
 
         if (doDensification || step >= stopSplitAt){
@@ -337,7 +339,6 @@ void Model::afterTrain(int step){
             int numPointsBefore = means.size(0);
 
             torch::Tensor culls = (torch::sigmoid(opacities) < cullAlphaThresh).squeeze();
-            int hugeCount = 0;
             if (splitsMask.numel()){
                 culls |= splitsMask;
             }
@@ -350,7 +351,6 @@ void Model::afterTrain(int step){
                     huge |= max2DSize > cullScreenSize;
                 }
                 culls |= huge;
-                hugeCount = torch::sum(huge).item<int>();
             }
 
             int cullCount = torch::sum(culls).item<int>();
@@ -369,6 +369,7 @@ void Model::afterTrain(int step){
                 removeFromOptimizer(featuresRestOpt, featuresRest, culls);
                 removeFromOptimizer(opacitiesOpt, opacities, culls);
                 
+                std::cout << "Culled " << (numPointsBefore - means.size(0)) << " gaussians, remaining " << means.size(0) << std::endl;
             }
         }
 
@@ -382,6 +383,7 @@ void Model::afterTrain(int step){
             auto paramState = std::make_unique<torch::optim::AdamParamState>(static_cast<torch::optim::AdamParamState&>(*opacitiesOpt->state()[pId]));
             paramState->exp_avg(torch::zeros_like(paramState->exp_avg()));
             paramState->exp_avg_sq(torch::zeros_like(paramState->exp_avg_sq()));
+            std::cout << "Alpha reset" << std::endl;
         }
 
         // Clear
diff --git a/opensplat.cpp b/opensplat.cpp
index 737741e..6ec7f53 100644
--- a/opensplat.cpp
+++ b/opensplat.cpp
@@ -94,6 +94,7 @@ int main(int argc, char *argv[]){
 
         InfiniteRandomIterator<ns::Camera> cams(inputData.cameras);
 
+        int imageSize = -1;
         for (size_t step = 1; step <= numIters; step++){
             ns::Camera cam = cams.next();
 
@@ -103,17 +104,22 @@ int main(int argc, char *argv[]){
             torch::Tensor gt = cam.getImage(model.getDownscaleFactor(step));
             gt = gt.to(device);
 
+            if (gt.size(0) != imageSize){
+                imageSize = gt.size(0) + 1;
+                std::cout << "Image size " << imageSize << "px" << std::endl;
+            }
+
             torch::Tensor ssimLoss = 1.0f - model.ssim.eval(rgb, gt);
             torch::Tensor l1Loss = ns::l1(rgb, gt);
             torch::Tensor mainLoss = (1.0f - ssimWeight) * l1Loss + ssimWeight * ssimLoss;
             mainLoss.backward();
+            
+            if (step % 10 == 0) std::cout << "Step " << step << ": " << mainLoss.item<float>() << std::endl;
 
             model.optimizersStep();
             //model.optimizersScheduleStep(); // TODO
             model.afterTrain(step);
 
-            if (step % 10 == 0) std::cout << "Step " << step << ": " << mainLoss.item<float>() << std::endl;
-
             if (saveEvery > 0 && step % saveEvery == 0){
                 fs::path p(outputScene);
                 model.savePlySplat(p.replace_filename(fs::path(p.stem().string() + "_" + std::to_string(step) + p.extension().string()).string()));