cameras_from_opencv_projection device #1021

Summary: Fix https://github.com/facebookresearch/pytorch3d/issues/1021 that cameras_from_opencv_projection always creates on CPU.

Reviewed By: nikhilaravi

Differential Revision: D33508211

fbshipit-source-id: fadebd45cacafd633af6a58094cf6f654529992c
This commit is contained in:
Jeremy Reizenstein 2022-01-21 05:29:04 -08:00 committed by Facebook GitHub Bot
parent 39bb2ce063
commit 45d096e219
2 changed files with 8 additions and 4 deletions

View File

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

View File

@ -68,7 +68,7 @@ class TestCameraConversions(TestCaseMixin, unittest.TestCase):
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 @@ class TestCameraConversions(TestCaseMixin, unittest.TestCase):
]
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 @@ class TestCameraConversions(TestCaseMixin, unittest.TestCase):
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(