train.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2026 Povl Filip Sonne-Frederiksen
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#pragma once
 
#include <reusex/gsplat/GaussianCloud.hpp>
#include <reusex/gsplat/TrainingViews.hpp>
 
#include <atomic>
#include <cstdint>
#include <filesystem>
#include <string>
#include <vector>
 
namespace cv {
class Mat;
}
 
namespace reusex {
class ProjectDB;
}
 
namespace reusex::gsplat {

struct MCMCOptions {
  bool enabled = false;

  double cap_factor = 2.0;
  std::int64_t cap_absolute = 0;

  double growth = 1.05;
 
  int refine_every = 100; 
  int refine_start = 500; 
  double refine_stop_fraction = 0.85;

  float min_opacity = 0.005f;
 
  float noise_lr = 5e5f;          
  float noise_opacity_t = 0.005f; 
  float noise_opacity_k = 100.0f; 

  float opacity_reg = 0.0f;
  float scale_reg = 0.0f;
};

struct TrainOptions {
  int iterations = 2000;
 
  float lr_means = 1.6e-4f;   
  float lr_scales = 5e-3f;    
  float lr_quats = 1e-3f;     
  float lr_opacities = 5e-2f; 
  float lr_sh_dc = 2.5e-3f;   

  float lr_sh_rest = 2.5e-3f / 20.0f;

  int sh_degree_interval = 1000;

  float lambda_dssim = 0.2f;

  bool prune_enabled = true;
  int prune_interval = 500;      
  int prune_start = 500;         
  float prune_opacity = 0.005f;  
  float prune_max_scale = 10.0f; 

  MCMCOptions mcmc;

  int holdout_every = 8;

  int eval_interval = 1000;

  int eval_max_views = 32;

  unsigned seed = 42;
 
  int log_interval = 100; 

  std::vector<int> render_iterations;
  std::filesystem::path render_dir;
  std::size_t render_view_index = 0;

  const std::atomic_bool *cancel_token = nullptr;

  int checkpoint_every = 0;
  int checkpoint_keep = 3;
  std::filesystem::path checkpoint_dir;
};

struct TrainMetrics {
  int iteration = 0;
  double loss = 0; 
  double l1 = 0;   
  double psnr = 0; 
  std::size_t gaussians = 0;
};

struct EvalMetrics {
  int iteration = 0;
  double train_psnr = 0; 
  double train_ssim = 0;
  double holdout_psnr = 0; 
  double holdout_ssim = 0;
  std::size_t train_views = 0;
  std::size_t holdout_views = 0;
  std::size_t gaussians = 0;
};
 
struct TrainResult {
  GaussianCloud gaussians;
  std::vector<TrainMetrics> history;
  std::vector<EvalMetrics> evals;
  double seconds = 0;
  std::size_t final_count = 0;
  double final_psnr = 0; 
  double final_holdout_psnr = 0;
  double final_holdout_ssim = 0;
  std::size_t relocated = 0; 
  std::size_t added = 0;     

  std::vector<std::size_t> trained_views;
  std::vector<std::size_t> holdout_view_indices;
  std::vector<std::filesystem::path> renders;

  std::vector<std::filesystem::path> checkpoints;

  int final_sh_degree = 0;

  bool cancelled = false;
  int iterations_run = 0;
};

bool is_available();

bool has_cuda_device();

TrainResult train_gaussians(const GaussianCloud &init,
                            const std::vector<TrainingView> &views,
                            const TrainOptions &opt = {});

cv::Mat render_view(const GaussianCloud &gaussians, const TrainingView &view);

struct GsplatStageOptions {
  std::string seed_cloud = "cloud"; 
  GaussianInitOptions init;
  TrainingViewOptions views;
  TrainOptions train;

  std::string splat_name = "splat";

  std::filesystem::path out_ply;
};

TrainResult run_gsplat_stage(ProjectDB &db, const GsplatStageOptions &opt);
 
} // namespace reusex::gsplat