From 34b1b4ab8bb2dd619e57d230cdaf8b0a35196a85 Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@fb.com>
Date: Sat, 16 Oct 2021 14:40:55 -0700
Subject: [PATCH] defaulted grid_sizes in points2vols

Summary: Fix #873, that grid_sizes defaults to the wrong dtype in points2volumes code, and mask doesn't have a proper default.

Reviewed By: nikhilaravi

Differential Revision: D31503545

fbshipit-source-id: fa32a1a6074fc7ac7bdb362edfb5e5839866a472
---
 pytorch3d/common/workaround/symeig3x3.py |  2 +-
 pytorch3d/ops/points_to_volumes.py       |  6 +++++-
 tests/bm_symeig3x3.py                    |  3 ++-
 tests/test_iou_box3d.py                  |  1 +
 tests/test_points_to_volumes.py          | 16 +++++++++++++++-
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/pytorch3d/common/workaround/symeig3x3.py b/pytorch3d/common/workaround/symeig3x3.py
index 52198d41..42467b98 100644
--- a/pytorch3d/common/workaround/symeig3x3.py
+++ b/pytorch3d/common/workaround/symeig3x3.py
@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from typing import Tuple, Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.nn.functional as F
diff --git a/pytorch3d/ops/points_to_volumes.py b/pytorch3d/ops/points_to_volumes.py
index 9c87a89b..0d18903b 100644
--- a/pytorch3d/ops/points_to_volumes.py
+++ b/pytorch3d/ops/points_to_volumes.py
@@ -364,7 +364,7 @@ def add_points_features_to_volume_densities_features(
         # grid sizes shape (minibatch, 3)
         grid_sizes = (
             torch.LongTensor(list(volume_densities.shape[2:]))
-            .to(volume_densities)
+            .to(volume_densities.device)
             .expand(volume_densities.shape[0], 3)
         )
 
@@ -386,6 +386,10 @@ def add_points_features_to_volume_densities_features(
         splat = False
     else:
         raise ValueError('No such interpolation mode "%s"' % mode)
+
+    if mask is None:
+        mask = points_3d.new_ones(1).expand(points_3d.shape[:2])
+
     volume_densities, volume_features = _points_to_volumes(
         points_3d,
         points_features,
diff --git a/tests/bm_symeig3x3.py b/tests/bm_symeig3x3.py
index 7a345462..bb51c66a 100644
--- a/tests/bm_symeig3x3.py
+++ b/tests/bm_symeig3x3.py
@@ -6,7 +6,7 @@
 
 
 from itertools import product
-from typing import Callable, Any
+from typing import Any, Callable
 
 import torch
 from common_testing import get_random_cuda_device
@@ -14,6 +14,7 @@ from fvcore.common.benchmark import benchmark
 from pytorch3d.common.workaround import symeig3x3
 from test_symeig3x3 import TestSymEig3x3
 
+
 torch.set_num_threads(1)
 
 CUDA_DEVICE = get_random_cuda_device()
diff --git a/tests/test_iou_box3d.py b/tests/test_iou_box3d.py
index f40ae3a3..990b1f45 100644
--- a/tests/test_iou_box3d.py
+++ b/tests/test_iou_box3d.py
@@ -16,6 +16,7 @@ from pytorch3d.io import save_obj
 from pytorch3d.ops.iou_box3d import _box_planes, _box_triangles, box3d_overlap
 from pytorch3d.transforms.rotation_conversions import random_rotation
 
+
 OBJECTRON_TO_PYTORCH3D_FACE_IDX = [0, 4, 6, 2, 1, 5, 7, 3]
 DATA_DIR = get_tests_dir() / "data"
 DEBUG = False
diff --git a/tests/test_points_to_volumes.py b/tests/test_points_to_volumes.py
index 0e190ae6..311d3330 100644
--- a/tests/test_points_to_volumes.py
+++ b/tests/test_points_to_volumes.py
@@ -12,7 +12,10 @@ from typing import Tuple
 import numpy as np
 import torch
 from common_testing import TestCaseMixin
-from pytorch3d.ops import add_pointclouds_to_volumes
+from pytorch3d.ops import (
+    add_pointclouds_to_volumes,
+    add_points_features_to_volume_densities_features,
+)
 from pytorch3d.ops.points_to_volumes import _points_to_volumes
 from pytorch3d.ops.sample_points_from_meshes import sample_points_from_meshes
 from pytorch3d.structures.meshes import Meshes
@@ -373,6 +376,17 @@ class TestPointsToVolumes(TestCaseMixin, unittest.TestCase):
                     else:
                         self.assertTrue(torch.isfinite(field.grad.data).all())
 
+    def test_defaulted_arguments(self):
+        points = torch.rand(30, 1000, 3)
+        features = torch.rand(30, 1000, 5)
+        _, densities = add_points_features_to_volume_densities_features(
+            points,
+            features,
+            torch.zeros(30, 1, 32, 32, 32),
+            torch.zeros(30, 5, 32, 32, 32),
+        )
+        self.assertClose(torch.sum(densities), torch.tensor(30 * 1000.0), atol=0.1)
+
     def _check_volume_slice_color_density(
         self, V, split_dim, interp_mode, clr_gt, slice_type, border=3
     ):