Fix focal lenght dtype + add support for different znear,zfar

Author: jkulhanek

internal/camera_utils.py

@@ -602,7 +602,7 @@ def pix_to_dir(x, y):
         dll = dl
       else:
         dll = dl[:, mask, :2]
-      theta = xnp.sqrt(xnp.sum(xnp.square(dll[..., :2]), axis=-1))
+      theta = xnp.sqrt(xnp.sum(xnp.square(dll[..., :2]), axis=-1, keepdims=True))
       theta = xnp.minimum(xnp.pi, theta)
       sin_theta_over_theta = xnp.sin(theta) / theta
 

internal/datasets.py

@@ -329,7 +329,11 @@ def __init__(self,
     if len(self.pixtocams.shape) == 2:
       self.pixtocams = np.repeat(self.pixtocams[None], self._n_examples, 0)
     if not isinstance(self.focal, np.ndarray):
-      self.focal = np.full((self._n_examples,), self.focal, dtype=np.int32)
+      self.focal = np.full((self._n_examples,), self.focal, dtype=np.float32)
+    if not isinstance(self.near, np.ndarray):
+      self.near = np.full((self._n_examples,), self.near, dtype=np.float32)
+    if not isinstance(self.height, np.ndarray):
+      self.far = np.full((self._n_examples,), self.far, dtype=np.float32)
     if not isinstance(self.width, np.ndarray):
       self.width = np.full((self._n_examples,), self.width, dtype=np.int32)
     if not isinstance(self.height, np.ndarray):
@@ -450,10 +454,11 @@ def _make_ray_batch(self,
     """
 
     broadcast_scalar = lambda x: np.broadcast_to(x, pix_x_int.shape)[..., None]
+    idx = 0 if self.render_path else cam_idx
     ray_kwargs = {
         'lossmult': broadcast_scalar(1.) if lossmult is None else lossmult,
-        'near': broadcast_scalar(self.near),
-        'far': broadcast_scalar(self.far),
+        'near': broadcast_scalar(self.near[idx]),
+        'far': broadcast_scalar(self.far[idx]),
         'cam_idx': broadcast_scalar(cam_idx),
     }
     # Collect per-camera information needed for each ray.
@@ -532,7 +537,11 @@ def generate_ray_batch(self, cam_idx: int) -> utils.Batch:
     if self._render_spherical:
       camtoworld = self.camtoworlds[cam_idx]
       rays = camera_utils.cast_spherical_rays(
-          camtoworld, self.height[cam_idx], self.width[cam_idx], self.near, self.far, xnp=np)
+          camtoworld,
+          self.height[cam_idx], self.width[cam_idx],
+          self.near[cam_idx],
+          self.far[cam_idx],
+          xnp=np)
       return utils.Batch(rays=rays)
     else:
       # Generate rays for all pixels in the image.