tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

bugfix: handle multiple camera types in a single batch (#3415)

Browse Source
This commit is contained in:
Daniel Crispell 2024-09-10 15:31:05 -04:00 committed by GitHub
parent d67b281f7d
commit e2a0915e56
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
nerfstudio/cameras/cameras.py
View File

@ -778,15 +778,15 @@ class Cameras(TensorDataclass):
return vr180_origins, directions_stack
for cam in cam_types:
if CameraType.PERSPECTIVE.value in cam_types:
for cam_type in cam_types:
if CameraType.PERSPECTIVE.value == cam_type:
mask = (self.camera_type[true_indices] == CameraType.PERSPECTIVE.value).squeeze(-1) # (num_rays)
mask = torch.stack([mask, mask, mask], dim=0)
directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0], mask).float()
directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1], mask).float()
directions_stack[..., 2][mask] = -1.0
elif CameraType.FISHEYE.value in cam_types:
elif CameraType.FISHEYE.value == cam_type:
mask = (self.camera_type[true_indices] == CameraType.FISHEYE.value).squeeze(-1) # (num_rays)
mask = torch.stack([mask, mask, mask], dim=0)
@ -803,7 +803,7 @@ class Cameras(TensorDataclass):
).float()
directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask).float()
elif CameraType.EQUIRECTANGULAR.value in cam_types:
elif CameraType.EQUIRECTANGULAR.value == cam_type:
mask = (self.camera_type[true_indices] == CameraType.EQUIRECTANGULAR.value).squeeze(-1) # (num_rays)
mask = torch.stack([mask, mask, mask], dim=0)
@ -816,22 +816,22 @@ class Cameras(TensorDataclass):
directions_stack[..., 1][mask] = torch.masked_select(torch.cos(phi), mask).float()
directions_stack[..., 2][mask] = torch.masked_select(-torch.cos(theta) * torch.sin(phi), mask).float()
elif CameraType.OMNIDIRECTIONALSTEREO_L.value in cam_types:
elif CameraType.OMNIDIRECTIONALSTEREO_L.value == cam_type:
ods_origins_circle, directions_stack = _compute_rays_for_omnidirectional_stereo("left")
# assign final camera origins
c2w[..., :3, 3] = ods_origins_circle
elif CameraType.OMNIDIRECTIONALSTEREO_R.value in cam_types:
elif CameraType.OMNIDIRECTIONALSTEREO_R.value == cam_type:
ods_origins_circle, directions_stack = _compute_rays_for_omnidirectional_stereo("right")
# assign final camera origins
c2w[..., :3, 3] = ods_origins_circle
elif CameraType.VR180_L.value in cam_types:
elif CameraType.VR180_L.value == cam_type:
vr180_origins, directions_stack = _compute_rays_for_vr180("left")
# assign final camera origins
c2w[..., :3, 3] = vr180_origins
elif CameraType.VR180_R.value in cam_types:
elif CameraType.VR180_R.value == cam_type:
vr180_origins, directions_stack = _compute_rays_for_vr180("right")
# assign final camera origins
c2w[..., :3, 3] = vr180_origins
@ -880,7 +880,7 @@ class Cameras(TensorDataclass):
directions_stack[coord_mask] = camera_utils.fisheye624_unproject(masked_coords, camera_params)
else:
raise ValueError(f"Camera type {cam} not supported.")
raise ValueError(f"Camera type {cam_type} not supported.")
assert directions_stack.shape == (3,) + num_rays_shape + (3,)