Make transform_points_screen's with_xyflip configurable

Krzysztof Chalupka · facebook-github-bot · commit 437200198186 · 2022-05-16T18:23:51.000-07:00
Summary: We'll need non-flipped screen coords in splatter.

Reviewed By: bottler

Differential Revision: D36337027

fbshipit-source-id: 897f88e8854bab215d2d0e502b25d15526ee86f1
diff --git a/pytorch3d/renderer/cameras.py b/pytorch3d/renderer/cameras.py
@@ -323,7 +323,7 @@ def transform_points_ndc(
         return world_to_ndc_transform.transform_points(points, eps=eps)
 
     def transform_points_screen(
-        self, points, eps: Optional[float] = None, **kwargs
+        self, points, eps: Optional[float] = None, with_xyflip: bool = True, **kwargs
     ) -> torch.Tensor:
         """
         Transforms points from PyTorch3D world/camera space to screen space.
@@ -341,14 +341,19 @@ def transform_points_screen(
                 stabilizes gradients since it leads to avoiding division
                 by excessively low numbers for points close to the
                 camera plane.
+            with_xyflip: If True, flip x and y directions. In world/camera/ndc coords,
+                +x points to the left and +y up. If with_xyflip is true, in screen
+                coords +x points right, and +y down, following the usual RGB image
+                convention. Warning: do not set to False unless you know what you're
+                doing!
 
         Returns
             new_points: transformed points with the same shape as the input.
         """
         points_ndc = self.transform_points_ndc(points, eps=eps, **kwargs)
         image_size = kwargs.get("image_size", self.get_image_size())
         return get_ndc_to_screen_transform(
-            self, with_xyflip=True, image_size=image_size
+            self, with_xyflip=with_xyflip, image_size=image_size
         ).transform_points(points_ndc, eps=eps)
 
     def clone(self):
diff --git a/tests/test_camera_pixels.py b/tests/test_camera_pixels.py
@@ -204,6 +204,9 @@ def test_camera(self):
         for cameras in (data.camera_ndc, data.camera_screen):
             ndc_points = cameras.transform_points_ndc(points)
             screen_points = cameras.transform_points_screen(points)
+            screen_points_without_xyflip = cameras.transform_points_screen(
+                points, with_xyflip=False
+            )
             camera_points = cameras.transform_points(points)
             for batch_idx in range(2):
                 # NDC space agrees with the original
@@ -214,6 +217,13 @@ def test_camera(self):
                     torch.tensor([data.x + 0.5, data.y + 0.5, 1.0]),
                     atol=1e-5,
                 )
+                # Screen coords without xyflip should have x, y that negate the non-
+                # flipped values, and unchanged z.
+                self.assertClose(
+                    screen_points_without_xyflip[batch_idx][0],
+                    torch.tensor([-(data.x + 0.5), -(data.y + 0.5), 1.0]),
+                    atol=1e-5,
+                )
                 # Second point in screen space is the center of the screen
                 self.assertClose(
                     screen_points[batch_idx][1],
