tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

adding cat to tensordataclass

Browse Source
This commit is contained in:
Ethan Weber 2023-08-04 15:47:39 +00:00
parent 7e90cb7989
commit f4be4017de
2 changed files with 34 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nerfstudio/scripts/render.py
View File

@ -497,6 +497,8 @@ class RenderInterpolated(BaseRender):
"""Frame rate of the output video."""
output_format: Literal["images", "video"] = "video"
"""How to save output data."""
loop: bool = False
"""Whether to start and end with the first camera."""
def main(self) -> None:
"""Main function."""
@ -515,6 +517,10 @@ class RenderInterpolated(BaseRender):
assert pipeline.datamanager.train_dataset is not None
cameras = pipeline.datamanager.train_dataset.cameras
if self.loop:
# This will make the camera path start and end with the same camera
cameras = cameras.cat(cameras[:1])
seconds = self.interpolation_steps * len(cameras) / self.frame_rate
camera_path = get_interpolated_camera_path(
cameras=cameras,

34
nerfstudio/utils/tensor_dataclass.py
View File

@ -145,12 +145,12 @@ class TensorDataclass:
def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT:
if isinstance(indices, (torch.Tensor)):
return self._apply_fn_to_fields(lambda x: x[indices])
return self._apply_fn_to_fields(lambda field_name, x: x[indices])
if isinstance(indices, (int, slice, type(Ellipsis))):
indices = (indices,)
assert isinstance(indices, tuple)
def tensor_fn(x):
def tensor_fn(field_name, x):
return x[indices + (slice(None),)]
def dataclass_fn(x):
@ -195,6 +195,28 @@ class TensorDataclass:
"""Returns the number of dimensions of the tensor dataclass."""
return len(self._shape)
def cat(self: TensorDataclassT, other: TensorDataclassT, dim: int = 0) -> TensorDataclassT:
"""Returns a new TensorDataclass with the other TensorDataclass's values concatenated along the specified dimension.
Args:
other: The other TensorDataclass to concatenate with.
dim: The dimension along which to concatenate the tensors.
Returns:
A new TensorDataclass with the tensors concatenated along the specified dimension.
"""
if not isinstance(other, TensorDataclass):
raise TypeError("Expected 'other' to be a TensorDataclass.")
def tensor_fn(field_name, x):
return torch.cat([x, getattr(other, field_name)], dim=dim)
dataclass_fn = None
custom_tensor_dims_fn = None
return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn)
def reshape(self: TensorDataclassT, shape: Tuple[int, ...]) -> TensorDataclassT:
"""Returns a new TensorDataclass with the same data but with a new shape.
@ -209,7 +231,7 @@ class TensorDataclass:
if isinstance(shape, int):
shape = (shape,)
def tensor_fn(x):
def tensor_fn(field_name, x):
return x.reshape((*shape, x.shape[-1]))
def dataclass_fn(x):
@ -247,7 +269,7 @@ class TensorDataclass:
return v.broadcast_to((*shape, *v.shape[-custom_dims:]))
return self._apply_fn_to_fields(
lambda x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn
lambda field_name, x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn
)
def to(self: TensorDataclassT, device) -> TensorDataclassT:
@ -259,7 +281,7 @@ class TensorDataclass:
Returns:
A new TensorDataclass with the same data but on the specified device.
"""
return self._apply_fn_to_fields(lambda x: x.to(device))
return self._apply_fn_to_fields(lambda field_name, x: x.to(device))
def _apply_fn_to_fields(
self: TensorDataclassT,
@ -330,7 +352,7 @@ class TensorDataclass:
):
new_dict[f] = custom_tensor_dims_fn(f, v)
elif isinstance(v, (torch.Tensor, TensorDataclass)):
new_dict[f] = fn(v)
new_dict[f] = fn(f, v)
elif isinstance(v, Dict):
new_dict[f] = self._apply_fn_to_dict(v, fn, dataclass_fn)
else: