diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index b28b60b90..a78d94ddc 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/nerfstudio/model_components/scene_colliders.py @@ -108,49 +108,44 @@ class AABBBoxCollider(SceneCollider): return ray_bundle +@torch.jit.script +def _intersect_with_sphere( + rays_o: torch.Tensor, rays_d: torch.Tensor, center: torch.Tensor, radius: float = 1.0, near_plane: float = 0.0 +): + a = (rays_d * rays_d).sum(dim=-1, keepdim=True) + b = 2 * (rays_o - center) * rays_d + b = b.sum(dim=-1, keepdim=True) + c = (rays_o - center) * (rays_o - center) + c = c.sum(dim=-1, keepdim=True) - radius**2 + + # clamp to near plane + nears = (-b - torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a) + fars = (-b + torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a) + + nears = torch.clamp(nears, min=near_plane) + fars = torch.maximum(fars, nears + 1e-6) + + nears = torch.nan_to_num(nears, nan=0.0) + fars = torch.nan_to_num(fars, nan=0.0) + + return nears, fars + + class SphereCollider(SceneCollider): """Module for colliding rays with the scene box to compute near and far values. Args: - center: center of sphere to intersect + center: center of sphere to intersect [3] redius: radius of sphere to intersect + near_plane: near plane to clamp to """ - def __init__(self, center: TensorType[3], radius: float, near_plane: float = 0.0, **kwargs) -> None: + def __init__(self, center: torch.Tensor, radius: float, near_plane: float = 0.0, **kwargs) -> None: super().__init__(**kwargs) self.center = center self.radius = radius self.near_plane = near_plane - def _intersect_with_sphere(self, rays_o: TensorType["num_rays":..., 3], rays_d: TensorType["num_rays":..., 3]): - """Returns collection of valid rays within a specified near/far bounding box along with a mask - specifying which rays are valid - - Args: - rays_o: (num_rays, 3) ray origins - rays_d: (num_rays, 3) ray directions - """ - - a = (rays_d * rays_d).sum(dim=-1, keepdim=True) - self.center = self.center.to(rays_o) - b = 2 * (rays_o - self.center) * rays_d - b = b.sum(dim=-1, keepdim=True) - c = (rays_o - self.center) * (rays_o - self.center) - c = c.sum(dim=-1, keepdim=True) - self.radius**2 - - # clamp to near plane - near_plane = self.near_plane if self.training else 0 - nears = (-b - torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a) - fars = (-b + torch.sqrt(torch.square(b) - 4 * a * c)) / (2 * a) - - nears = torch.clamp(nears, min=near_plane) - fars = torch.maximum(fars, nears + 1e-6) - - nears = torch.nan_to_num(nears, nan=0.0) - fars = torch.nan_to_num(fars, nan=0.0) - - return nears, fars - def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle: """Intersects the rays with the scene box and updates the near and far values. Populates nears and fars fields and returns the ray_bundle. @@ -158,7 +153,15 @@ class SphereCollider(SceneCollider): Args: ray_bundle: specified ray bundle to operate on """ - nears, fars = self._intersect_with_sphere(ray_bundle.origins, ray_bundle.directions) + self.center = self.center.to(ray_bundle.origins.device) + near_plane = self.near_plane if self.training else 0 + nears, fars = _intersect_with_sphere( + rays_o=ray_bundle.origins, + rays_d=ray_bundle.directions, + center=self.center, + radius=self.radius, + near_plane=near_plane, + ) ray_bundle.nears = nears ray_bundle.fars = fars return ray_bundle diff --git a/pyproject.toml b/pyproject.toml index 42fd7b146..7c4249edf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ "opencv-python==4.6.0.66", "Pillow>=9.3.0", "plotly>=5.7.0", - "protobuf<=3.20.3.*,!=3.20.0", + "protobuf<=3.20.3", "pymeshlab>=2022.2.post2", "pyngrok>=5.1.0", "python-socketio>=5.7.1", diff --git a/scripts/texture.py b/scripts/texture.py index ec6545b87..702242ef2 100644 --- a/scripts/texture.py +++ b/scripts/texture.py @@ -56,10 +56,10 @@ class TextureMesh: # texture the mesh with NeRF and export to a mesh.obj file # and a material and texture file texture_utils.export_textured_mesh( - mesh, - pipeline, - self.px_per_uv_triangle, - self.output_dir, + mesh=mesh, + pipeline=pipeline, + output_dir=self.output_dir, + px_per_uv_triangle=self.px_per_uv_triangle, unwrap_method=self.unwrap_method, num_pixels_per_side=self.num_pixels_per_side, )