Export Splatfacto-optimized poses

[oseiskar]

Currently, ns-export cameras does not export the optimized poses even if Splatfacto pose optimization is enabled. Similarly, ns-eval does not take the pose optimization into account. This is because the "adjustment" to the initial pose is stored as a separate variable in the camera optimizer, which allows properly using the initial poses as priors. The NeRF code uses a different approach.

This PR is a slightly cleaned-up version of the approach used in https://github.com/SpectacularAI/3dgs-deblur, which is basically ensuring that the "camera index" and an instance of the CameraOptimizer are available in the relevant places. However, this seems a bit ugly and fragile.

Any ideas of to improve this? Permanently "applying" the adjustments to the camera_to_worlds at the end of training could be one option, which would probably help with ns-export cameras (but not the eval PSNR metrics exported to, e.g., tensorboard, which are also fixed in this PR).

This is also somewhat related to #3073.

Testing

First, download suitable data and install this branch (similarly to here):

conda activate nerfstudio

# download a suitable dataset
git clone https://github.com/SpectacularAI/3dgs-deblur
cd 3dgs-deblur
python download_data.py --dataset synthetic
cd ..

# check out and install this branch
https://github.com/SpectacularAI/nerfstudio
cd nerfstudio
git checkout export-splatfacto-optimized-poses
pip install -e .

Train (example):

ns-train splatfacto --data ../3dgs-deblur/data/inputs-processed/synthetic-posenoise/factory \
    --pipeline.model.camera-optimizer.mode=SO3xR3 \
    --pipeline.model.rasterize-mode antialiased \
    --max-num-iterations 15000 \
    --output-dir outputs/pose-opt-test \
    nerfstudio-data --eval-mode all

Check metrics:

ns-eval --load-config outputs/pose-opt-test/splatfacto/TIMESTAMP/config.yml

Then see the metrics in output.json. With this PR, should show high PSNR (> 32). Without this, shows something quite low (< 20) due to the misaligned evaluation poses.

In addition, the poses exported with the following command are different (with this PR, they are the optimized ones, otherwise the original ones, possibly after some auto-scaling/transformations by the dataparser):

ns-export cameras \
    --load-config outputs/pose-opt-test/splatfacto/TIMESTAMP/config.yml \
    --output-dir outputs/pose-opt-test/

[ichsan2895]

I see.. When I activate SO3xR3, I checked the image eval result, the rendered image is little bit shifting due to camera optimization vs GT. I think it was cuprit for metrics (PSNR, LPIPS, SSIM) degradation. So, it corrected the eval images aligned with corrected poses.

BTW, @jb-ye has overhauled the splatfacto in this draft (#3113). Hopefully it wont break each other.

[jb-ye]

I am hesitant about exporting the optimized eval poses. This is very much only for the evaluation metrics reporting for paper write-ups. In general, I'm afraid we don't have a proper framework design in nerfstudio to handle eval pose optimization.

I think it would be useful to export both original training poses and optimized training poses. At least this is useful to inspect the results. So just make this PR focus on training poses, no? What do you think?

[oseiskar]

I see.. When I activate SO3xR3, I checked the image eval result, the rendered image is little bit shifting due to camera optimization vs GT. I think it was cuprit for metrics (PSNR, LPIPS, SSIM) degradation. So, it corrected the eval images aligned with corrected poses.

@ichsan2895 Sorry, I did not fully get this. Did you mean that this happens in the current Splatfacto or with this PR. What should happen is that in the current main branch, if you activate SO3xR3, you get the misalignment and metric degradation, but this PR fixes the issue.

I am hesitant about exporting the optimized eval poses. This is very much only for the evaluation metrics reporting for paper write-ups. In general, I'm afraid we don't have a proper framework design in nerfstudio to handle eval pose optimization.

@jb-ye I agree. Currently, the proper framework for optimizing eval poses (i.e., optimizing their poses while keeping all Gaussian parameters fixed) is missing. For this reason, I have only tested the --eval-mode=all option (other --eval-modes are broken with pose optimization anyway). Then training and evaluation sets are the same. However, maybe also in that case, it would be a cleaner solution to only output the optimized training poses and leave the evaluation poses as the originals.

I think it would be useful to export both original training poses and optimized training poses. At least this is useful to inspect the results. So just make this PR focus on training poses, no? What do you think?

Currently, both the original and optmized training poses are available in the saved model. I think exporting the optimized training poses is a more reasonable default for ns-export cameras, but it could also be possible to add a command line flag there that would also allow exporting the original poses.

The original poses are also presumably available in the input data, but they do not necessarily match the output of ns-export cameras due to the various transforms in data parsers, so I also agree that both unoptimized and optimized training poses are potentially useful export targets.

[jb-ye]

it would be a cleaner solution to only output the optimized training poses and leave the evaluation poses as the originals.

Yes, I think keep the PR simple is important for long term maintenance.

I think exporting the optimized training poses is a more reasonable default for ns-export cameras, but it could also be possible to add a command line flag there that would also allow exporting the original poses.

That could also work. Export the original poses (post various transforms), so they can be used to compare with the optimized pose.

[oseiskar]

Update:

• rebased on latest main
• fixed linter warnings
• exporting optimized poses only for training images

To summarize, the following aspects should now be fixed in Splatfacto pose optimization with the options --pipeline.model.camera-optimizer.mode=SO3xR3 --eval-mode all:

• camera poses saved to transforms_train.json with ns-export cameras
• metric output from ns-eval (EDIT: removed due to fragile implementation)
• "Eval Images Metrics" graphs in Tensorboard

The following are left as future improvements:

• Correct handling of other eval modes (this is a bigger revamp, as mentioned here)
• Command line parameter for exporting unoptimized poses in ns-export cameras
• metric output from ns-eval
• "Eval Images Metrics" graphs in Tensorboard

Also note that the --eval-mode all case is relevant for practical reconstruction cases where there is not separate test/evaluation set.

[jb-ye]

This is a custom code for SplatfactoModel. But I also have Nerfacto model which have camera optimizer enabled by default. Can we support both?

[oseiskar]

Hmm. Good point. In theory, yes, but if we want a generic solution instead of separate ifs for Spaltfacto and Nerfacto, we would also need to fix this: #3073

[jb-ye]

Maybe we should address #3073 first, and use the same API call for apply_to_camera()? The additional change seems very manageable in my opinion.

keep the current behavior for NeRFs, but match the splatfacto behavior for consistency

[oseiskar]

Ok. I'll get back to this a bit later

[oseiskar]

@jb-ye Done: I tested that I now can:

1. optimize poses using either Splatfacto or Nerfacto (with --eval-mode all and --pipeline.model.camera-optimizer.mode=SO3xR3)
2. ns-export cameras
3. modify the input data by replacing the original noisy poses with the optimized poses
4. train Splatfacto on the modified input data without pose optimization, and get good results

My scripts for step 3 are a bit messy (based on this, mostly), but at least in principle, it now works.

[jb-ye]

Based on #3073 's discussion, the apply_to_raybundle is neither multiplying transform from left or right. I believe, one should return

rotations = torch.bmm(adj[..., :3, :3], camera.camera_to_worlds[..., :3, :3])
translations = camera.camera_to_worlds[..., :3, 3] + adj[..., :3, 3]

[oseiskar]

Thanks for checking. You're right... Quite weird that this actually worked quite well in my test. Perhaps the angular adjustments were so small in the synthetic test that the difference between the two formulas was negligible. I'll anyway change it to match...

[oseiskar]

@jb-ye Now actually fixed. Here's my test results that show the subtle difference

1. Artificial pose noise, Splatfacto reconstruction without pose optimization
2. Poses optimized with Nerfacto, reconstruction with Splatfacto without pose opt. old version
3. Same as above but with the new version

1

2

3

[jb-ye]

Thanks for the effort! Looks good to me.