Add utils to approximate the conical frustums as multivariate gaussians.

Summary:
Introduce methods to approximate the radii of conical frustums along rays as described in [MipNerf](https://arxiv.org/abs/2103.13415):

- Two new attributes are added to ImplicitronRayBundle: bins and radii. Bins is of size n_pts_per_ray + 1. It allows us to manipulate easily and n_pts_per_ray intervals. For example we need the intervals coordinates in the radii computation for \(t_{\mu}, t_{\delta}\). Radii are used to store the radii of the conical frustums.

- Add 3 new methods to compute the radii:
   - approximate_conical_frustum_as_gaussians: It computes the mean along the ray direction, the variance of the
      conical frustum  with respect to t and variance of the conical frustum with respect to its radius. This
      implementation follows the stable computation defined in the paper.
   - compute_3d_diagonal_covariance_gaussian: Will leverage the two previously computed variances to find the
     diagonal covariance of the Gaussian.
   - conical_frustum_to_gaussian: Mix everything together to compute the means and the diagonal covariances along
     the ray of the Gaussians.

- In AbstractMaskRaySampler, introduces the attribute `cast_ray_bundle_as_cone`. If False it won't change the previous behaviour of the RaySampler. However if True, the samplers will sample `n_pts_per_ray +1` instead of `n_pts_per_ray`. This points are then used to set the bins attribute of ImplicitronRayBundle. The support of HeterogeneousRayBundle has not been added since the current code does not allow it. A safeguard has been added to avoid a silent bug in the future.

Reviewed By: shapovalov

Differential Revision: D45269190

fbshipit-source-id: bf22fad12d71d55392f054e3f680013aa0d59b78

pytorch3d/implicitron/models/utils.py

@@ -177,6 +177,20 @@ def chunk_generator(
 
     for start_idx in iter:
         end_idx = min(start_idx + chunk_size_in_rays, n_rays)
+        bins = (
+            None
+            if ray_bundle.bins is None
+            else ray_bundle.bins.reshape(batch_size, n_rays, n_pts_per_ray + 1)[
+                :, start_idx:end_idx
+            ]
+        )
+        pixel_radii_2d = (
+            None
+            if ray_bundle.pixel_radii_2d is None
+            else ray_bundle.pixel_radii_2d.reshape(batch_size, -1, 1)[
+                :, start_idx:end_idx
+            ]
+        )
         ray_bundle_chunk = ImplicitronRayBundle(
             origins=ray_bundle.origins.reshape(batch_size, -1, 3)[:, start_idx:end_idx],
             directions=ray_bundle.directions.reshape(batch_size, -1, 3)[
@@ -186,6 +200,8 @@ def chunk_generator(
                 :, start_idx:end_idx
             ],
             xys=ray_bundle.xys.reshape(batch_size, -1, 2)[:, start_idx:end_idx],
+            bins=bins,
+            pixel_radii_2d=pixel_radii_2d,
             camera_ids=_safe_slice(ray_bundle.camera_ids, start_idx, end_idx),
             camera_counts=_safe_slice(ray_bundle.camera_counts, start_idx, end_idx),
         )