Support splat export in original dataset coordinates

[brentyi]

The .ply export for Gaussians previously defaulted to saving in Nerfstudio's auto-oriented / auto-scaled coordinate frame.

I added support for exporting this in the original dataset coordinate frame, which the point cloud export supports via a "Save in world frame" checkbox.

I matched this in both the GUI and the export script CLI:

To me it makes the most sense to default this to True (I also flipped this for the point cloud export), but open to thoughts!

cc @jb-ye, #2909

[kerrj]

seems good

[jb-ye]

I don't think the option would work actually for spherical harmonic parameters. The spherical harmonics are saved in the transformed frame and cannot be "re-orient" back to the original coordinates. (We can apply translation and scale, but not rotations, otherwise SH becomes inconsistent.)

My recommendation is to save the transforms as meta data as part of exporter.

[brentyi]

I don't think the option would work actually for spherical harmonic parameters. The spherical harmonics are saved in the transformed frame and cannot be "re-orient" back to the original coordinates. (We can apply translation and scale, but not rotations, otherwise SH becomes inconsistent.)

My recommendation is to save the transforms as meta data as part of exporter.

Agree the current implementation is flawed, will revisit this after ECCV deadline!

See comments about SH parameters

[pwais]

@jb-ye Many viewers ignore spherical harmonics, and in my experience doing a rigid transform on splats looks plenty fine. I think ideally Nerfstudio just does a scale and re-center internally as part of the model rather than changing any of the camera poses, initial points etc that are part of the dataset. It should be up to the model to normalize as needed, not up to the user.

[pwais]

please make this configurable? or maybe this is just for debugging

[brentyi]

I'm actually not sure why this shows up in this diff, the change is from #2969. it speeds up undistortion a lot for big datasets I've been toying with!

it's hardcoded to 2 because we can really only expect benefits from one thread doing IO while the other thread is doing undistortion; the implementation is still weird given this (ideally we'd just have 1 worker doing IO while the main one is sequentially undistorting) but I'm a fan of not letting perfect be the enemy of... better 🤷

[pwais]

FWIW each worker might need cv2.setNumThreads(1) or else it can choke the CPU. I seem to see way more than 200% util here hence why i brought it up so maybe it's just not tuned well for all users.

maybe in a future refactor this stuff will just get pushed to a torch dataloader... the pinned memory part breaks for me for large datasets anyways, literally I got a OOM and hard lock-up because too too too much much much pinned memory

[pwais]

pretty please don't do transform math w/out at least comments, this sort of code is 110% likely to put a future reader in transform hell

also pretty please use pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space() because
(1) that's what's used elsewhere in this file
(2) using that function ensures future refactors won't break things, and most past nerfstudio refactors have indeed broken lots of things

[pwais]

please don't do this, ESPECIALLY w/out comments. I have lost a lot of time reading nerfstudio code that's like this. instead consider:

positions = model.means.cpu().numpy()

poses = np.eye(4, dtype=np.float32)[None, ...].repeat(positions.shape[0], axis=0)[:, :3, :]
poses[:, :3, 3] = positions
poses = pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space(
    torch.from_numpy(poses)
)

[pwais]

??? First of all, scipy.spatial.transform already has quaternion multiply... at least this code is clear about scalar-first versus scalar-last.

this could be made more concise, but consider instead:

from scipy.spatial.transform import Rotation as ScR
# ns gplat says quaternions are [w,x,y,z] scalar-first format
# scipy is [x, y, z, w] scalar-last format

raw_quats = model.quats.data.cpu().numpy().squeeze()
R_quats = ScR.from_quat(raw_quats[:, [1, 2, 3, 0]])

# apply the inverse dataparser transform to the splat rotations
poses = np.eye(4, dtype=np.float32)[None, ...].repeat(raw_quats.shape[0], axis=0)[:, :3, :]
poses[:, :3, :3] = R_quats.as_matrix()
poses = pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space(
    torch.from_numpy(poses)
)
rots_in_input = poses[:, :3, :3].numpy()
quat_in_input = ScR.from_matrix(rots_in_input)
quats = quat_in_input.as_quat()[:, [3, 0, 1, 2], None]

Again, this uses transform_poses_to_original_space(), which might amalgamate several different transforms and scales, who knows? Instead of trying to re-derive the transform as the current PR does. And hopefully transform_poses_to_original_space() gets maintained. But it's really really important to be clear about frames etc, and transform_poses_to_original_space() helps with that a ton.

[brentyi]

(note for when we revive this PR, which is planned) for the quaternion multiply if we don't want to deal with the xyzw/wxyz conversion of scipy we can also use (vtf.SO3(wxyz0) @ vtf.SO3(wxyz1)).wxyz with import viser.transforms as vtf where viser>=0.1.30

[pwais]

voicing a preference for the use of standard scipy / numpy / torch wherever possible

yes it's unfortunate that there are different quaternion encodings, different camera conventions, different euler angle conventions ....

[brentyi]

I think ideally Nerfstudio just does a scale and re-center internally as part of the model rather than changing any of the camera poses, initial points etc that are part of the dataset.

Tough to overcome the inertia here but I agree that this would solve a lot of problems!

[pwais]

Tough to overcome the inertia here but I agree that this would solve a lot of problems!

Yes inertia but at least the PR discussions leave breadcrumbs. Things that fall off the rolling katamari ball become seeds for the next re-write.

[jkulhanek]

Wait, spherical harmonics can easily be rotated by using Wigner matrices, right?

[jb-ye]

@jkulhanek Could you share a gist of sample code of rotating spherical harmonics? I am not aware of Wigner matrices.

[jkulhanek]

Just a note here. Since SH is a complete basis, the rotation is possible exactly. Here is the code:
https://github.com/jkulhanek/nerfbaselines/blob/develop/nerfbaselines/math_utils.py

The code is heavy, and I don't understand it fully, but I played with it in a notebook, and it seems to rotate the SH correctly. I can also share the notebook if you want.

It's the recursive impl (whatever it means) which is supposed to be more stable (for higher order). The good thing about it is that the wigner matrix can be computed once and then applied to all SHs at once.

[Frenchman997]

Is there any update on this? I would also be interested in the ability to export splats in their original world frame. I haven't looked into it yet at all but would be willing to contribute.

[hardikdava]

@jkulhanek can you provide the notebook? Thanks

[hongsukchoi]

Spherical harmonics functions are equivariant to SO(3) and can be rotated.

I don't know the exact code to do it, but probably possible with e3nn api:

From ChatGPT

import torch
from e3nn.o3 import Irrep, Irreps
from e3nn.o3 import spherical_harmonics, Rotation

R = torch.tensor([ [0.36, 0.48, -0.8], [-0.8, 0.60, 0], [0.48, 0.64, 0.6] ])

l = 2  # Degree of spherical harmonics

directions = torch.tensor([[0.0, 0.0, 1.0]]) # z-axis unit vector 

Y_lm = spherical_harmonics(l, directions)

# Get the irreducible representation (Irrep) of the spherical harmonics of degree 
l irrep = Irrep(f"{l}e") # e denotes even parity 

# Apply the rotation to the spherical harmonics 
Y_lm_rotated = irrep.D_from_matrix(R) @ Y_lm

Explanation:

• spherical_harmonics(l, directions): Computes the spherical harmonics of degree l for the given directions.
• Irrep(f"{l}e"): Creates an irreducible representation for the spherical harmonics of degree l. The "e" stands for even parity, which is typical for standard spherical harmonics.
• D_from_matrix(R): Computes the Wigner D-matrix for the rotation matrix R, which describes how the spherical harmonics transform under rotation.
• @: Performs matrix multiplication between the Wigner D-matrix and the original spherical harmonics to obtain the rotated spherical harmonics.

[Ben-Mack]

Is Playcanvas related code for SH rotation a good fit to unblock this PR? https://github.com/playcanvas/engine/blob/release-1.69/src/scene/gsplat/gsplat.js#L259

Also, how can we get the matrix that represent the output coordinate? (Or the diff between original and output coordinate)

[pwais]

The line you cite seems to just load splats? rather than rotate the SHs given an arbitrary transform.

Also, how can we get the matrix that represent the output coordinate? (Or the diff between original and output coordinate)

In nerfstudio there's a dataparser_transform that gets stored in the dataparser outputs / model metadata files. AFAIK there's no "splat file format" that normally contains this transform. The code in this change uses that transform to put the trained splats back into the training frame (right now the PR just has translations and rotations tho I believe, yeah).

[jkulhanek]

You can use this code https://gist.github.com/jkulhanek/aae3ec12d779ffc729c72157315df0da

[Ben-Mack]

Very helpful explanation, thank you!

The line you cite seems to just load splats? rather than rotate the SHs given an arbitrary transform.

Sorry, this is the correct link, which they mentioned Rotate spherical harmonics up to band 3 based on https://github.com/andrewwillmott/sh-lib:
https://github.com/playcanvas/supersplat/blob/b008eca2bd28d2aae7d2defd67363ae6c4f7174b/src/sh-utils.ts#L42

[pwais]

on first blush it's hard to compare the cited PlayCanvas impl vs @jkulhanek 's link (PlayCanvas looks possibly like an algebraic simplification somehow?).

But it would be nice for any final solution to have some unit tests. For example, if dataparser_transform breaks somehow, want to easily determine it's that instead of something in complex sh rotations.