Merge branch 'main' into ethan/benchmarks

nerfstudio/model_components/scene_colliders.py

@@ -108,6 +108,65 @@ 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 [3]
+        redius: radius of sphere to intersect
+        near_plane: near plane to clamp to
+    """
+
+    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 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.
+
+        Args:
+            ray_bundle: specified ray bundle to operate on
+        """
+        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
+
+
 class NearFarCollider(SceneCollider):
     """Sets the nears and fars with fixed values.
 

pyproject.toml

@@ -37,7 +37,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,
         )