From 61dc520e17fb38beefaf2fb7288c4355ec353fb7 Mon Sep 17 00:00:00 2001
From: Matthew Tancik <tancik@berkeley.edu>
Date: Mon, 30 Jan 2023 10:55:13 -0800
Subject: [PATCH] Add sphere collider (#1310)

Co-authored-by: David McAllister <davidrmc2@gmail.com>
---
 .../model_components/scene_colliders.py       | 59 +++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py
index bf3ccd1b9..a36edfb38 100644
--- a/nerfstudio/model_components/scene_colliders.py
+++ b/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.