Merge branch 'main' into generative

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

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",

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,
         )