Add BlenderCamera

Summary: Adding BlenderCamera (for rendering with R2N2 Blender transformations in the next diff).

Reviewed By: nikhilaravi

Differential Revision: D22462515

fbshipit-source-id: 4b40ee9bba8b6d56788dd3c723036ec704668153

Author: megluyagao

pytorch3d/datasets/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
-from .r2n2 import R2N2
+from .r2n2 import R2N2, BlenderCamera
 from .shapenet import ShapeNetCore
 from .utils import collate_batched_meshes
 

pytorch3d/datasets/r2n2/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
-from .r2n2 import R2N2
+from .r2n2 import R2N2, BlenderCamera
 
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

pytorch3d/datasets/r2n2/r2n2.py

@@ -11,11 +11,18 @@
 from PIL import Image
 from pytorch3d.datasets.shapenet_base import ShapeNetBase
 from pytorch3d.io import load_obj
+from pytorch3d.renderer.cameras import CamerasBase
+from pytorch3d.transforms import Transform3d
 from tabulate import tabulate
 
 
 SYNSET_DICT_DIR = Path(__file__).resolve().parent
 
+# Default values of rotation, translation and intrinsic matrices for BlenderCamera.
+r = np.expand_dims(np.eye(3), axis=0)  # (1, 3, 3)
+t = np.expand_dims(np.zeros(3), axis=0)  # (1, 3)
+k = np.expand_dims(np.eye(4), axis=0)  # (1, 4, 4)
+
 
 class R2N2(ShapeNetBase):
     """
@@ -217,3 +224,27 @@ def __getitem__(self, model_idx, view_idxs: Optional[List[int]] = None) -> Dict:
             model["images"] = torch.stack(images)
 
         return model
+
+
+class BlenderCamera(CamerasBase):
+    """
+    Camera for rendering objects with calibration matrices from the R2N2 dataset
+    (which uses Blender for rendering the views for each model).
+    """
+
+    def __init__(self, R=r, T=t, K=k, device="cpu"):
+        """
+        Args:
+            R: Rotation matrix of shape (N, 3, 3).
+            T: Translation matrix of shape (N, 3).
+            K: Intrinsic matrix of shape (N, 4, 4).
+            device: torch.device or str.
+        """
+        # The initializer formats all inputs to torch tensors and broadcasts
+        # all the inputs to have the same batch dimension where necessary.
+        super().__init__(device=device, R=R, T=T, K=K)
+
+    def get_projection_transform(self, **kwargs) -> Transform3d:
+        transform = Transform3d(device=self.device)
+        transform._matrix = self.K.transpose(1, 2).contiguous()  # pyre-ignore[16]
+        return transform

tests/test_r2n2.py

@@ -11,13 +11,16 @@
 import torch
 from common_testing import TestCaseMixin, load_rgb_image
 from PIL import Image
-from pytorch3d.datasets import R2N2, collate_batched_meshes
+from pytorch3d.datasets import R2N2, BlenderCamera, collate_batched_meshes
 from pytorch3d.renderer import (
     OpenGLPerspectiveCameras,
     PointLights,
     RasterizationSettings,
     look_at_view_transform,
 )
+from pytorch3d.renderer.cameras import get_world_to_view_transform
+from pytorch3d.transforms import Transform3d
+from pytorch3d.transforms.so3 import so3_exponential_map
 from torch.utils.data import DataLoader
 
 
@@ -258,3 +261,21 @@ def test_render_r2n2(self):
                 "test_r2n2_render_by_categories_%s.png" % idx, DATA_DIR
             )
             self.assertClose(mixed_rgb, image_ref, atol=0.05)
+
+    def test_blender_camera(self):
+        """
+        Test BlenderCamera.
+        """
+        # Test get_world_to_view_transform.
+        T = torch.randn(10, 3)
+        R = so3_exponential_map(torch.randn(10, 3) * 3.0)
+        RT = get_world_to_view_transform(R=R, T=T)
+        cam = BlenderCamera(R=R, T=T)
+        RT_class = cam.get_world_to_view_transform()
+        self.assertTrue(torch.allclose(RT.get_matrix(), RT_class.get_matrix()))
+        self.assertTrue(isinstance(RT, Transform3d))
+
+        # Test getting camera center.
+        C = cam.get_camera_center()
+        C_ = -torch.bmm(R, T[:, :, None])[:, :, 0]
+        self.assertTrue(torch.allclose(C, C_, atol=1e-05))