Support splat export in original dataset coordinates #2951
Conversation
brentyi
changed the title
|
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! |
…dio-project/nerfstudio into brent/splat_export_world_frame
There was a problem hiding this comment.
@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.
| ) | ||
| ) | ||
|
|
||
| CONSOLE.log("Caching / undistorting eval images") | ||
| with ThreadPoolExecutor() as executor: | ||
| with ThreadPoolExecutor(max_workers=2) as executor: |
There was a problem hiding this comment.
please make this configurable? or maybe this is just for debugging
There was a problem hiding this comment.
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 🤷
There was a problem hiding this comment.
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
|
|
||
| output_scale = 1 / dataparser_scale | ||
| output_transform = np.zeros((3, 4)) | ||
| output_transform[:3, :3] = dataparser_transform[:3, :3].T |
There was a problem hiding this comment.
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
| np.einsum("ij,bj->bi", output_transform[:3, :3], model.means.cpu().numpy() * output_scale) | ||
| + output_transform[None, :3, 3] |
There was a problem hiding this comment.
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)
)
| for i in range(3): | ||
| map_to_tensors[f"scale_{i}"] = scales[:, i, None] | ||
|
|
||
| quats = model.quats.data.cpu().numpy() | ||
| def quaternion_multiply(wxyz0: np.ndarray, wxyz1: np.ndarray) -> np.ndarray: |
There was a problem hiding this comment.
??? 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.
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. |
|
Wait, spherical harmonics can easily be rotated by using Wigner matrices, right? |
|
@jkulhanek Could you share a gist of sample code of rotating spherical harmonics? I am not aware of Wigner matrices. |
|
Just a note here. Since SH is a complete basis, the rotation is possible exactly. Here is the code: 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. |
|
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. |
|
@jkulhanek can you provide the notebook? Thanks |
The
.plyexport 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