gs/config.py

"""
Configuration settings and constants for 3D Gaussian Splatting with NeRF datasets.
"""
import warp as wp
import numpy as np
import random

SEED = 42
random.seed(SEED)

# Warp data types and constants (keep capitalized as they are types)
WP_FLOAT16 = wp.float16
WP_FLOAT32 = wp.float32
WP_INT = wp.int32
WP_VEC2 = wp.vec2
WP_VEC2H = wp.vec2h
VEC6 = wp.types.vector(length=6, dtype=WP_FLOAT32)
DEVICE = "cuda" #"cpu" # Use "cpu" or "cuda"

TILE_M = wp.constant(16)
TILE_N = wp.constant(16)
TILE_THREADS = wp.constant(256)


class GaussianParams:
    """Parameters for 3D Gaussian Splatting."""

    # Training parameters
    num_iterations = 3*7000//1  # Default number of training iterations
    num_points = 5000 # Initial number of Gaussian points

    # Simple learning rate scheduler configuration
    use_lr_scheduler = True 
    # Learning rate scheduler configuration
    lr_scheduler_config = {
        'lr_pos': 1e-4,      # Initial learning rate for positions
        'lr_scale': 5e-4,    # Initial learning rate for scales  
        'lr_rot': 5e-3,      # Initial learning rate for rotations
        'lr_sh': 2e-3,       # Initial learning rate for spherical harmonics
        'lr_opac': 5e-3,     # Initial learning rate for opacities
        'final_lr_factor': 0.01  # Final LR will be 1% of initial LR
    }

    # Optimization parameters
    densification_interval = 999  # Perform densification every N iterations
    pruning_interval = 99  # Perform pruning every N iterations
    opacity_reset_interval = 5001
    save_interval = 99  # Save checkpoint every N iterations
    adam_beta1 = 0.9  # Adam optimizer beta1 parameter
    adam_beta2 = 0.999  # Adam optimizer beta2 parameter
    adam_epsilon = 1e-8  # Adam optimizer epsilon parameter
    
    
    densify_grad_threshold = 0.0002
    cull_opacity_threshold = 0.005
    start_prune_iter = 500
    densify_until_iter = 15000  # Stop densification after this iteration

    percent_dense = 0.01
    max_allowed_prune_ratio = 0.1 # no limit on pruning ratio

    # Gaussian parameters
    initial_scale = 0.0035 #0.1  # Initial scale for Gaussian points
    scale_modifier = 1.0  # Scaling factor for Gaussian splats
    sh_degree = 1 # Spherical harmonics degree

    # Scene parameters
    scene_scale = 1.0  # Scale factor for the scene
    background_color = [1.0,1.0,1.0] #[0.0, 0.0, 0.0]  # White background for NeRF synthetic

    # Loss parameters
    lambda_dssim = 0.2  # Weight for SSIM loss (1.0 means only SSIM, 0.0 means only L1)
    depth_loss_weight = 0.5  # Weight for depth L1 loss
    
    near = 0.01  # Default near clipping plane
    far = 100.0  # Default far clipping plane

    @classmethod
    def update(cls, **kwargs):
        """Update parameters with new values."""
        for key, value in kwargs.items():
            if hasattr(cls, key):
                setattr(cls, key, value)
            else:
                raise ValueError(f"Unknown parameter: {key}")

    """Configuration for Dynamic 4D Gaussian Splatting."""
    
    # Training iterations
    initial_frame_iterations = 3000     # Iterations for frame 0 (canonical)
    subsequent_frame_iterations = 500   # Iterations for frames 1, 2, ...
    
    # Dynamic parameter flags
    # By default, only positions vary over time
    # Set these to True to enable time-varying versions
    dynamic_positions = True            # Positions always vary (4D tracking)
    dynamic_scales = False              # Scales fixed across time
    dynamic_rotations = True           # Rotations fixed across time
    dynamic_opacities = False           # Opacities fixed across time
    dynamic_shs = False                 # SH coefficients fixed across time
    
    # 2DGS specific settings
    isotropic_gaussians = False          # If True, sx=sy (circular disks instead of ellipses)
    
    # Point cloud processing
    conf_threshold = 50.0               # Confidence percentile threshold
    min_points = 1000                   # Minimum points to keep
    max_points = 500000                 # Maximum points (for memory)
    
    # Optimization
    position_lr = 0.001                 # Learning rate for positions
    shared_param_lr_multiplier = 1.0    # Multiplier for shared param LRs
    
    # Sequential training specific
    copy_positions_from_previous = True # Initialize frame t from frame t-1
    
    # Rendering
    background_color = [1.0, 1.0, 1.0]  # White background
    
    @classmethod
    def get_config_dict(cls):
        """Get full configuration as dictionary (includes both static and dynamic params)."""
        return {
            # Original GaussianParams
            'num_iterations': cls.num_iterations,
            'num_points': cls.num_points,
            'densification_interval': cls.densification_interval,
            'pruning_interval': cls.pruning_interval,
            'scale_modifier': cls.scale_modifier,
            'sh_degree': cls.sh_degree,
            'background_color': cls.background_color,
            'save_interval': cls.save_interval,
            'adam_beta1': cls.adam_beta1,
            'adam_beta2': cls.adam_beta2,
            'adam_epsilon': cls.adam_epsilon,
            'initial_scale': cls.initial_scale,
            'scene_scale': cls.scene_scale,
            'near': cls.near,
            'far': cls.far,
            'lambda_dssim': cls.lambda_dssim,
            'depth_loss_weight': cls.depth_loss_weight,
            'densify_grad_threshold': cls.densify_grad_threshold,
            'cull_opacity_threshold': cls.cull_opacity_threshold,
            'start_prune_iter': cls.start_prune_iter,
            'densify_until_iter': cls.densify_until_iter,
        
            'use_lr_scheduler': cls.use_lr_scheduler,
            'lr_scheduler_config': cls.lr_scheduler_config,
            'max_allowed_prune_ratio': cls.max_allowed_prune_ratio,
            # Dynamic 4D Gaussian params
            'initial_frame_iterations': cls.initial_frame_iterations,
            'subsequent_frame_iterations': cls.subsequent_frame_iterations,
            'dynamic_positions': cls.dynamic_positions,
            'dynamic_scales': cls.dynamic_scales,
            'dynamic_rotations': cls.dynamic_rotations,
            'dynamic_opacities': cls.dynamic_opacities,
            'dynamic_shs': cls.dynamic_shs,
            'isotropic_gaussians': cls.isotropic_gaussians,
            'conf_threshold': cls.conf_threshold,
            'min_points': cls.min_points,
            'max_points': cls.max_points,
            'position_lr': cls.position_lr,
            'shared_param_lr_multiplier': cls.shared_param_lr_multiplier,
            'copy_positions_from_previous': cls.copy_positions_from_previous,
        }
    
    @classmethod
    def update(cls, **kwargs):
        """Update parameters with new values."""
        for key, value in kwargs.items():
            if hasattr(cls, key):
                setattr(cls, key, value)
            else:
                raise ValueError(f"Unknown parameter: {key}")