-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support splat export in original dataset coordinates #2951
base: main
Are you sure you want to change the base?
Changes from all commits
a24faf6
cc999fc
b86d7e1
f9ae770
55639be
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -32,8 +32,10 @@ | |
import tyro | ||
from typing_extensions import Annotated, Literal | ||
|
||
from nerfstudio.cameras.camera_utils import quaternion_from_matrix | ||
from nerfstudio.cameras.rays import RayBundle | ||
from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager | ||
from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanager | ||
from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager | ||
from nerfstudio.data.scene_box import OrientedBox | ||
from nerfstudio.exporter import texture_utils, tsdf_utils | ||
|
@@ -121,7 +123,7 @@ class ExportPointCloud(Exporter): | |
"""Number of rays to evaluate per batch. Decrease if you run out of memory.""" | ||
std_ratio: float = 10.0 | ||
"""Threshold based on STD of the average distances across the point cloud to remove outliers.""" | ||
save_world_frame: bool = False | ||
save_world_frame: bool = True | ||
"""If set, saves the point cloud in the same frame as the original dataset. Otherwise, uses the | ||
scaled and reoriented coordinate space expected by the NeRF models.""" | ||
|
||
|
@@ -482,6 +484,11 @@ class ExportGaussianSplat(Exporter): | |
Export 3D Gaussian Splatting model to a .ply | ||
""" | ||
|
||
save_world_frame: bool = True | ||
"""If set, saves the splat in the same frame as the original dataset. | ||
Otherwise, uses the scaled and reoriented coordinate space produced | ||
internally by Nerfstudio.""" | ||
|
||
def main(self) -> None: | ||
if not self.output_dir.exists(): | ||
self.output_dir.mkdir(parents=True) | ||
|
@@ -497,7 +504,26 @@ def main(self) -> None: | |
map_to_tensors = {} | ||
|
||
with torch.no_grad(): | ||
positions = model.means.cpu().numpy() | ||
if self.save_world_frame: | ||
assert isinstance(pipeline.datamanager, FullImageDatamanager) | ||
dataparser_outputs = pipeline.datamanager.train_dataparser_outputs | ||
dataparser_scale = dataparser_outputs.dataparser_scale | ||
dataparser_transform = dataparser_outputs.dataparser_transform.numpy(force=True) | ||
|
||
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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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 |
||
output_transform[:3, 3] = -dataparser_transform[:3, :3].T @ dataparser_transform[:3, 3] | ||
else: | ||
output_scale = 1 | ||
output_transform = np.zeros((3, 4)) | ||
output_transform[:3, :3] = np.eye(3) | ||
inv_dataparser_quat = quaternion_from_matrix(output_transform[:3, :3]) | ||
|
||
positions = ( | ||
np.einsum("ij,bj->bi", output_transform[:3, :3], model.means.cpu().numpy() * output_scale) | ||
+ output_transform[None, :3, 3] | ||
Comment on lines
+524
to
+525
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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:
|
||
) | ||
n = positions.shape[0] | ||
map_to_tensors["positions"] = positions | ||
map_to_tensors["normals"] = np.zeros_like(positions, dtype=np.float32) | ||
|
@@ -518,11 +544,27 @@ def main(self) -> None: | |
|
||
map_to_tensors["opacity"] = model.opacities.data.cpu().numpy() | ||
|
||
scales = model.scales.data.cpu().numpy() | ||
# Note that scales are in log space! | ||
scales = model.scales.data.cpu().numpy() + np.log(output_scale) | ||
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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ??? First of all, this could be made more concise, but consider instead:
Again, this uses There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. (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 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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 .... |
||
assert wxyz0.shape[-1] == 4 | ||
assert wxyz1.shape[-1] == 4 | ||
w0, x0, y0, z0 = np.moveaxis(wxyz0, -1, 0) | ||
w1, x1, y1, z1 = np.moveaxis(wxyz1, -1, 0) | ||
return np.stack( | ||
[ | ||
-x0 * x1 - y0 * y1 - z0 * z1 + w0 * w1, | ||
x0 * w1 + y0 * z1 - z0 * y1 + w0 * x1, | ||
-x0 * z1 + y0 * w1 + z0 * x1 + w0 * y1, | ||
x0 * y1 - y0 * x1 + z0 * w1 + w0 * z1, | ||
], | ||
axis=-1, | ||
) | ||
|
||
quats = quaternion_multiply(inv_dataparser_quat, model.quats.data.cpu().numpy()) | ||
for i in range(4): | ||
map_to_tensors[f"rot_{i}"] = quats[:, i, None] | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please make this configurable? or maybe this is just for debugging
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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