Points2vols doc fixes

Summary: Fixes to a couple of comments on points to volumes, make the mask work in round_points_to_volumes, and remove a duplicate rand calculation

Reviewed By: nikhilaravi

Differential Revision: D29522845

fbshipit-source-id: 86770ba37ef3942b909baf63fd73eed1399635b6

pytorch3d/ops/points_to_volumes.py

@@ -157,12 +157,15 @@ def add_points_features_to_volume_densities_features(
                 of its floating point coordinate. The weights are
                 determined using a trilinear interpolation scheme.
                 Trilinear splatting is fully differentiable w.r.t. all input arguments.
-        mask: A binary mask of shape `(minibatch, N)` determining which 3D points
-            are going to be converted to the resulting volume.
-            Set to `None` if all points are valid.
         min_weight: A scalar controlling the lowest possible total per-voxel
             weight used to normalize the features accumulated in a voxel.
             Only active for `mode==trilinear`.
+        mask: A binary mask of shape `(minibatch, N)` determining which 3D points
+            are going to be converted to the resulting volume.
+            Set to `None` if all points are valid.
+        grid_sizes: `LongTensor` of shape (minibatch, 3) representing the
+            spatial resolutions of each of the the non-flattened `volumes` tensors,
+            or None to indicate the whole volume is used for every batch element.
     Returns:
         volume_features: Output volume of shape `(minibatch, feature_dim, D, H, W)`
         volume_densities: Occupancy volume of shape `(minibatch, 1, D, H, W)`
@@ -284,13 +287,15 @@ def splat_points_to_volumes(
         volume_features: Batch of input *flattened* feature volumes
             of shape `(minibatch, feature_dim, N_voxels)`
         volume_densities: Batch of input *flattened* feature volume densities
-            of shape `(minibatch, 1, N_voxels)`. Each voxel should
+            of shape `(minibatch, N_voxels, 1)`. Each voxel should
             contain a non-negative number corresponding to its
             opaqueness (the higher, the less transparent).
         grid_sizes: `LongTensor` of shape (minibatch, 3) representing the
             spatial resolutions of each of the the non-flattened `volumes` tensors.
             Note that the following has to hold:
                 `torch.prod(grid_sizes, dim=1)==N_voxels`
+        min_weight: A scalar controlling the lowest possible total per-voxel
+            weight used to normalize the features accumulated in a voxel.
         mask: A binary mask of shape `(minibatch, N)` determining which 3D points
             are going to be converted to the resulting volume.
             Set to `None` if all points are valid.
@@ -457,9 +462,6 @@ def round_points_to_volumes(
     # split into separate coordinate vectors
     X, Y, Z = XYZ.split(1, dim=2)
 
-    # get random indices for the purpose of adding out-of-bounds values
-    rand_idx = X.new_zeros(X.shape).random_(0, n_voxels)
-
     # valid - binary indicators of votes that fall into the volume
     grid_sizes = grid_sizes.type_as(XYZ)
     valid = (
@@ -470,6 +472,8 @@ def round_points_to_volumes(
         * (0 <= Z)
         * (Z < grid_sizes_xyz[:, None, 2:3])
     ).long()
+    if mask is not None:
+        valid = valid * mask[:, :, None].long()
 
     # get random indices for the purpose of adding out-of-bounds values
     rand_idx = valid.new_zeros(X.shape).random_(0, n_voxels)

tests/test_points_to_volumes.py

@@ -30,8 +30,8 @@ def init_cube_point_cloud(
     batch_size: int = 10, n_points: int = 100000, rotate_y: bool = True
 ):
     """
-    Generate a random point cloud of `n_points` whose points of
-    which are sampled from faces of a 3D cube.
+    Generate a random point cloud of `n_points` whose points
+    are sampled from faces of a 3D cube.
     """
 
     # create the cube mesh batch_size times