From 4372001981868657ea2111575dc8704404a25090 Mon Sep 17 00:00:00 2001
From: Krzysztof Chalupka <chalupka@fb.com>
Date: Mon, 16 May 2022 18:23:51 -0700
Subject: [PATCH] Make transform_points_screen's with_xyflip configurable

Summary: We'll need non-flipped screen coords in splatter.

Reviewed By: bottler

Differential Revision: D36337027

fbshipit-source-id: 897f88e8854bab215d2d0e502b25d15526ee86f1
---
 pytorch3d/renderer/cameras.py |  9 +++++++--
 tests/test_camera_pixels.py   | 10 ++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/pytorch3d/renderer/cameras.py b/pytorch3d/renderer/cameras.py
index 002c21ec..d2d2d3cf 100644
--- a/pytorch3d/renderer/cameras.py
+++ b/pytorch3d/renderer/cameras.py
@@ -323,7 +323,7 @@ class CamerasBase(TensorProperties):
         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,6 +341,11 @@ class CamerasBase(TensorProperties):
                 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.
@@ -348,7 +353,7 @@ class CamerasBase(TensorProperties):
         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
index 72feaeb2..082246b6 100644
--- a/tests/test_camera_pixels.py
+++ b/tests/test_camera_pixels.py
@@ -204,6 +204,9 @@ class TestPixels(TestCaseMixin, unittest.TestCase):
         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 @@ class TestPixels(TestCaseMixin, unittest.TestCase):
                     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],