Bug fix for case where aspect ratio is a float

Summary:
- Fix the calculation of the non square NDC range when the H and W are not integer multiples.
- Add test for this case

Reviewed By: gkioxari

Differential Revision: D26613213

fbshipit-source-id: df6763cac602e9f1d516b41b432c4d2cfbaa356d

Author: nikhilaravi

pytorch3d/csrc/rasterize_points/rasterization_utils.cuh

@@ -10,7 +10,9 @@
 __device__ inline float NonSquareNdcRange(int S1, int S2) {
   float range = 2.0f;
   if (S1 > S2) {
-    range = ((S1 / S2) * range);
+    // First multiply S1 by float range so that division results
+    // in a float value.
+    range = (S1 * range) / S2;
   }
   return range;
 }

pytorch3d/csrc/rasterize_points/rasterization_utils.h

@@ -10,7 +10,7 @@
 inline float NonSquareNdcRange(int S1, int S2) {
   float range = 2.0f;
   if (S1 > S2) {
-    range = ((S1 / S2) * range);
+    range = (S1 * range) / S2;
   }
   return range;
 }

pytorch3d/renderer/points/rasterize_points.py

@@ -2,6 +2,7 @@
 
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 
 # pyre-fixme[21]: Could not find name `_C` in `pytorch3d`.
@@ -120,15 +121,7 @@ def rasterize_points(
             # Binned CPU rasterization not fully implemented
             bin_size = 0
         else:
-            # TODO: These heuristics are not well-thought out!
-            if max_image_size <= 64:
-                bin_size = 8
-            elif max_image_size <= 256:
-                bin_size = 16
-            elif max_image_size <= 512:
-                bin_size = 32
-            elif max_image_size <= 1024:
-                bin_size = 64
+            bin_size = int(2 ** max(np.ceil(np.log2(max_image_size)) - 4, 4))
 
     if bin_size != 0:
         # There is a limit on the number of points per bin in the cuda kernel.

tests/test_rasterize_rectangle_images.py

@@ -314,7 +314,7 @@ def _compare_square_with_nonsq(
 
             # Finally check the gradients of the input vertices for
             # the square and non square case
-            self.assertClose(verts_square.grad, grad_tensor.grad, rtol=2e-4)
+            self.assertClose(verts_square.grad, grad_tensor.grad, rtol=3e-4)
 
     def test_gpu(self):
         """
@@ -323,8 +323,9 @@ def test_gpu(self):
         dists, zbuf, bary are all the same for the square
         region which is present in both images.
         """
-        # Test both cases: (W > H), (H > W)
-        image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128)]
+        # Test both cases: (W > H), (H > W) as well as the case where
+        # H and W are not integer multiples of each other (i.e. float aspect ratio)
+        image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128), (600, 1110)]
 
         devices = ["cuda:0"]
         blurs = [0.0, 0.001]
@@ -391,7 +392,7 @@ def test_cpu(self):
         """
         # Test both when (W > H) and (H > W).
         # Using smaller image sizes here as the Python rasterizer is really slow.
-        image_sizes = [(32, 64), (64, 32)]
+        image_sizes = [(32, 64), (64, 32), (60, 110)]
         devices = ["cpu"]
         blurs = [0.0, 0.001]
         batch_sizes = [1]
@@ -646,8 +647,9 @@ def test_gpu(self):
         dists, zbuf, idx are all the same for the square
         region which is present in both images.
         """
-        # Test both cases: (W > H), (H > W)
-        image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128)]
+        # Test both cases: (W > H), (H > W) as well as the case where
+        # H and W are not integer multiples of each other (i.e. float aspect ratio)
+        image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128), (600, 1110)]
 
         devices = ["cuda:0"]
         blurs = [5e-2]
@@ -713,7 +715,7 @@ def test_cpu(self):
         """
         # Test both when (W > H) and (H > W).
         # Using smaller image sizes here as the Python rasterizer is really slow.
-        image_sizes = [(32, 64), (64, 32)]
+        image_sizes = [(32, 64), (64, 32), (60, 110)]
         devices = ["cpu"]
         blurs = [5e-2]
         batch_sizes = [1]