use gsplat strategy interface to simplify splatfacto

[jb-ye]

Given the latest updates of gsplat strategy interface, it is of interest to use the feature to benefit from multiple strategy implementations provided. The first step is to refactor existing splatfacto code meanwhile maintain the performance of current splafacto-big.

I spent some time debugging current gsplat implementation and found that we need to first merge this fix in order to reproduce the performance of current splatfacto-big. We don't necessarily need to merge the fix, this is just for producing baselines.

To benchmark simple_trainer.py in splatfacto configuration:

# using latest commit of gsplat and run in splatfacto-big equivalent configurations
python simple_trainer.py default --data_dir bicycle/ \
   --strategy.absgrad \
   --strategy.grow_grad2d 0.0006 \
   --strategy.prune_scale3d 0.5 \
   --strategy.refine_scale2d_stop_iter 4000 \
   --strategy.pause_refine_after_reset 300

The above setting in simple_trainer is "almost" equivalent to splatfacto-big (post this PR merge) with two differences:

• splatfacto-big have coarse-to-fine training enabled by default. Without coarse-to-fine training, absgrad tends to create a lot more gaussians in early iterations, which might not be necessarily useful.
• scene_scale calculation in nerfstudio colmap parser use abs(cam_positions) instead of 1.1 * norm(cam_positions). This uniformly affects all--strategy.*3d parameters.

Mipnerf360 bicycle

• Baseline A (original Inria):

  • 7k iterations: PSNR=23.59, SSIM=0.662, GS=3.57M
  • 30k iterations: PSNR=25.19, SSIM=0.763, GS=6.06M

• Baseline B (splatfacto-big in main):

  • 30k iterations: PSNR=25.25, SSIM=0.781, GS=5.36M

• simple_trainer (using splatfacto configuration) + fix:

  • 7k iterations: PSNR=23.96, SSIM=0.714, GS=5.06M
  • 30k iterations: PSNR=25.22, SSIM=0.780, GS=6.46M

• splatfacto-big (this PR) + fix:

  • 7k iterations: GS=3.9M
  • 30k iterations: PSNR=25.28, SSIM=0.781, GS=6.03M

• splatfacto-big (this PR, densify_grad_thresh=0.0006):

  • 7k iterations: GS=3.0M
  • 30k iterations: PSNR=25.24, SSIM=0.777, GS=4.70M

• splatfacto-big (this PR, densify_grad_thresh=0.0005):

  • 7k iterations: GS=3.8M
  • 30k iterations: PSNR=25.26, SSIM=0.781, GS=6.00M

Note that the default strategy is still a strong candidate if one chooses a proper configuration. The original Inria configuration is suboptimal.

[kerrj]

Looks great! thanks for taking this on. Are you planning on adding support for MCMC? it would be nice to remove hyperparameters from SplatfactoConfig and replace them with a sort of StrategyConfig which can be swapped out for different strategies. This might require a gsplat PR to format them in a compatible way with tyro.

I see that nerfstudio-project/gsplat#358 was closed without merging, is that change reflected inside the splatfacto hyperparams for gradient threhsolds now?

[jb-ye]

Are you planning on adding support for MCMC?

Yes, but I am also open if someone else wish to take on this task. This PR is part of my own learning about the new strategy interface. This is indeed very convenient.

it would be nice to remove hyperparameters from SplatfactoConfig and replace them with a sort of StrategyConfig which can be swapped out for different strategies. This might require a gsplat PR to format them in a compatible way with tyro.

Yes, it takes some thoughts about how to change the SplatfactoConfig.

I see that nerfstudio-project/gsplat#358 was closed without merging, is that change reflected inside the splatfacto hyperparams for gradient threhsolds now?

Yes, I retuned the parameter, I believe they are now very close.

[jb-ye]

Since this PR removed the --continue_cull_post_densification option, we will need to filter gaussians whose opacities are below 1/255. upon exporting.

[kerrj]

It seems num_sh_bases is no longer in the most recent version gsplat, we can fix this import error by manually defining it like this:

def num_sh_bases(degree: int) -> int:
    """
    Returns the number of spherical harmonic bases for a given degree.
    """
    return (degree + 1) ** 2

[jb-ye]

@kerrj Did we add a new test for building nerfstudio dockers? Not sure how to make it pass.

[brentyi]

@jkulhanek for Docker action, any idea?

[jkulhanek]

I think the error is that the docker build ran from a forked repo, but fork doesn’t have permission to run it. We should only run the docker action on the main repo - disable running it on forked repos I think.

[jkulhanek]

I've added a PR which should fix it.

[jb-ye]

@brentyi @kerrj
I think this PR is good to go.

[brentyi]

PR looks reasonable to me!

As a thought for supporting MCMC the easiest way I can think of for adding this is just putting into the config:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: DefaultStrategy = dataclasses.field(default_factory=DefaultStrategy)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

Or, if preventing breaking changes is more important than CLI flag organization:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: Annotated[DefaultStrategy, tyro.conf.arg(name="")] = dataclasses.field(
    default_factory=DefaultStrategy
)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

and then in the splatfacto constructor:

self.strategy = {
    "default": self.config.default_strategy,
    "mcmc": self.config.mcmc_strategy,
}[self.config.strategy_mode]

[abrahamezzeddine]

PR looks reasonable to me!

As a thought for supporting MCMC the easiest way I can think of for adding this is just putting into the config:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: DefaultStrategy = dataclasses.field(default_factory=DefaultStrategy)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

Or, if preventing breaking changes is more important than CLI flag organization:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: Annotated[DefaultStrategy, tyro.conf.arg(name="")] = dataclasses.field(
    default_factory=DefaultStrategy
)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

and then in the splatfacto constructor:

self.strategy = {
    "default": self.config.default_strategy,
    "mcmc": self.config.mcmc_strategy,
}[self.config.strategy_mode]

Hello,

Thank you for your awesome work.

I assume MCMC is not yet implemented. Do you have any update for when we could try this with nerfstudio?

[jb-ye]

@abrahamezzeddine someone created a PR for mcmc. But I haven't gotten time to look at it.

[Sunnyhong0326]

Hi, in this line, scene_scale is hard-coded to 1

nerfstudio/nerfstudio/models/splatfacto.py

Line 332 in 516fd7c

self.strategy_state = self.strategy.initialize_state(scene_scale=1.0)

From my understanding, scene_scale is the maximum of distance between average camera position and all camera positions.
When I trace the code of nerfstudio's colmap dataparser, the default setting is to auto scale all the camera poses and normalized the camera positions to -1 and 1. This results in scene_scale = 1.
However, in my application, I don't want to scale camera positions to -1 and 1. This results in scene_scale does not always = 1.

[jb-ye]

@Sunnyhong0326 this is a known limitation. Ideally, we should pass scene scale from parser to the model in nerfstudio framework. Let me see what I can do here