cameras_from_opencv_projection device #1021

Summary: Fix #1021 that cameras_from_opencv_projection always creates on CPU.

Reviewed By: nikhilaravi

Differential Revision: D33508211

fbshipit-source-id: fadebd45cacafd633af6a58094cf6f654529992c

Author: bottler

pytorch3d/renderer/camera_conversions.py

@@ -55,6 +55,7 @@ def _cameras_from_opencv_projection(
         focal_length=focal_pytorch3d,
         principal_point=p0_pytorch3d,
         image_size=image_size,
+        device=R.device,
     )
 
 

tests/test_camera_conversions.py

@@ -68,7 +68,7 @@ def test_opencv_conversion(self):
         return correct projections of random 3D points. The check is done
         against a set of results precomuted using `cv2.projectPoints` function.
         """
-
+        device = torch.device("cuda:0")
         image_size = [[480, 640]] * 4
         R = [
             [
@@ -116,17 +116,19 @@ def test_opencv_conversion(self):
         ]
 
         principal_point, focal_length, R, tvec, image_size = [
-            torch.FloatTensor(x)
+            torch.tensor(x, device=device)
             for x in (principal_point, focal_length, R, tvec, image_size)
         ]
-        camera_matrix = eyes(dim=3, N=4)
+        camera_matrix = eyes(dim=3, N=4, device=device)
         camera_matrix[:, 0, 0], camera_matrix[:, 1, 1] = (
             focal_length[:, 0],
             focal_length[:, 1],
         )
         camera_matrix[:, :2, 2] = principal_point
 
-        pts = torch.nn.functional.normalize(torch.randn(4, 1000, 3), dim=-1)
+        pts = torch.nn.functional.normalize(
+            torch.randn(4, 1000, 3, device=device), dim=-1
+        )
 
         # project the 3D points with the opencv projection function
         rvec = so3_log_map(R)
@@ -136,6 +138,7 @@ def test_opencv_conversion(self):
         cameras_opencv_to_pytorch3d = cameras_from_opencv_projection(
             R, tvec, camera_matrix, image_size
         )
+        self.assertEqual(cameras_opencv_to_pytorch3d.device, device)
 
         # project the 3D points with converted cameras to screen space.
         pts_proj_pytorch3d_screen = cameras_opencv_to_pytorch3d.transform_points_screen(