[Add] Get K, R, T from CameraParameter by a single method (#60)

* get krt and testr

* replace K_R_T with KRT

mmhuman3d/core/cameras/camera_parameters.py

@@ -117,11 +117,11 @@ def get_opencv_distort_mat(self) -> np.ndarray:
         dist_coeffs = np.array(dist_coeffs)
         return dist_coeffs
 
-    def set_K_R_T(self,
-                  K_mat: np.ndarray,
-                  R_mat: np.ndarray,
-                  T_vec: np.ndarray,
-                  inverse_extrinsic: bool = False) -> None:
+    def set_KRT(self,
+                K_mat: np.ndarray,
+                R_mat: np.ndarray,
+                T_vec: np.ndarray,
+                inverse_extrinsic: bool = False) -> None:
         """Set intrinsic and extrinsic of a camera.
 
         Args:
@@ -149,6 +149,42 @@ def set_K_R_T(self,
         self.set_mat_np('rotation_mat', R_mat)
         self.set_value('translation', T_vec.tolist())
 
+    def get_KRT(self, k_dim=3) -> List[np.ndarray]:
+        """Get intrinsic and extrinsic of a camera.
+
+        Args:
+            k_dim (int, optional):
+                Dimension of the returned mat K.
+                Defaults to 3.
+
+        Raises:
+            ValueError: k_dim is neither 3 nor 4.
+
+        Returns:
+            List[np.ndarray]:
+                K_mat (np.ndarray):
+                    In shape [3, 3].
+                R_mat (np.ndarray):
+                    Rotation from world to view in default.
+                    In shape [3, 3].
+                T_vec (np.ndarray):
+                    Translation from world to view in default.
+                    In shape [3,].
+        """
+        K_3x3 = self.get_mat_np('in_mat')
+        R_mat = self.get_mat_np('rotation_mat')
+        T_vec = np.asarray(self.get_value('translation'))
+        if k_dim == 3:
+            return [K_3x3, R_mat, T_vec]
+        elif k_dim == 4:
+            K_3x3 = np.expand_dims(K_3x3, 0)  # shape (1, 3, 3)
+            K_4x4 = convert_K_3x3_to_4x4(
+                K=K_3x3, is_perspective=True)  # shape (1, 4, 4)
+            K_4x4 = K_4x4[0, :, :]
+            return [K_4x4, R_mat, T_vec]
+        else:
+            raise ValueError(f'K mat cannot be converted to {k_dim}x{k_dim}')
+
     def set_mat_np(self, mat_key: str, mat_numpy: np.ndarray) -> None:
         """Set a matrix-type parameter to mat_numpy.
 
@@ -378,14 +414,9 @@ def export_to_perspective_cameras(self) -> PerspectiveCameras:
         """
         height = self.parameters_dict['H']
         width = self.parameters_dict['W']
-        k_3x3 = self.get_mat_np('in_mat')  # shape (3, 3)
-        k_3x3 = np.expand_dims(k_3x3, 0)  # shape (1, 3, 3)
-        k_4x4 = convert_K_3x3_to_4x4(
-            K=k_3x3, is_perspective=True)  # shape (1, 4, 4)
-        rotation = self.get_mat_np('rotation_mat')  # shape (3, 3)
+        k_4x4, rotation, translation = self.get_KRT(k_dim=4)
+        k_4x4 = np.expand_dims(k_4x4, 0)  # shape (1, 3, 3)
         rotation = np.expand_dims(rotation, 0)  # shape (1, 3, 3)
-        translation = self.get_value('translation')  # list, len==3
-        translation = np.asarray(translation)
         translation = np.expand_dims(translation, 0)  # shape (1, 3)
         new_K = torch.from_numpy(k_4x4)
         new_R = torch.from_numpy(rotation)

mmhuman3d/data/data_converters/h36m.py

@@ -133,7 +133,7 @@ def _get_camera_params(self, camera: int, subject: str) -> dict:
 
         camera_name = f'S{subject}_{self.camera_ids[camera]}'
         camera_params = CameraParameter(camera_name, h, w)
-        camera_params.set_K_R_T(K, R, T)
+        camera_params.set_KRT(K, R, T)
         camera_params.set_value('k1', float(k[0]))
         camera_params.set_value('k2', float(k[1]))
         camera_params.set_value('k3', float(k[2]))

mmhuman3d/data/data_converters/pw3d.py

@@ -118,7 +118,7 @@ def convert_by_mode(self, dataset_path: str, out_path: str,
                         T = extrinsic_param[:3, 3]
 
                         camera = CameraParameter(H=h, W=w)
-                        camera.set_K_R_T(K, R, T)
+                        camera.set_KRT(K, R, T)
                         parameter_dict = camera.to_dict()
 
                         pose[:3] = cv2.Rodrigues(

mmhuman3d/data/data_converters/surreal.py

@@ -161,7 +161,7 @@ def convert_by_mode(self, dataset_path: str, out_path: str,
                 _, R, T = self.get_extrinsic(annotations['camLoc'])
                 camera = CameraParameter(
                     H=self.image_height, W=self.image_width)
-                camera.set_K_R_T(K, R, T)
+                camera.set_KRT(K, R, T)
                 parameter_dict = camera.to_dict()
 
                 # image folder

tests/test_cameras/test_camera_parameter.py

@@ -18,13 +18,13 @@ def test_set_from_mat():
     mat_3x3 = np.eye(3)
     mat_4x4 = np.eye(4)
     vec_3 = np.zeros(shape=[3])
-    empty_param.set_K_R_T(K_mat=mat_3x3, R_mat=mat_3x3, T_vec=vec_3)
-    empty_param.set_K_R_T(
+    empty_param.set_KRT(K_mat=mat_3x3, R_mat=mat_3x3, T_vec=vec_3)
+    empty_param.set_KRT(
         K_mat=mat_3x3, R_mat=mat_3x3, T_vec=vec_3, inverse_extrinsic=True)
     with pytest.raises(AssertionError):
-        empty_param.set_K_R_T(K_mat=mat_4x4, R_mat=mat_3x3, T_vec=vec_3)
+        empty_param.set_KRT(K_mat=mat_4x4, R_mat=mat_3x3, T_vec=vec_3)
     with pytest.raises(AssertionError):
-        empty_param.set_K_R_T(K_mat=mat_3x3, R_mat=mat_4x4, T_vec=vec_3)
+        empty_param.set_KRT(K_mat=mat_3x3, R_mat=mat_4x4, T_vec=vec_3)
     assert len(empty_param.get_value('translation')) == 3
 
 
@@ -81,7 +81,7 @@ def test_misc():
     empty_param = CameraParameter(name='test_misc')
     mat_3x3 = np.eye(3)
     vec_3 = np.zeros(shape=[3])
-    empty_param.set_K_R_T(K_mat=mat_3x3, R_mat=mat_3x3, T_vec=vec_3)
+    empty_param.set_KRT(K_mat=mat_3x3, R_mat=mat_3x3, T_vec=vec_3)
     empty_param.set_value('k1', 1.0)
     empty_param.reset_distort()
     distort_vec = empty_param.get_opencv_distort_mat()
@@ -103,3 +103,11 @@ def test_misc():
         empty_param.get_value('distortion_k1')
     dumped_str = empty_param.to_string()
     assert isinstance(dumped_str, str)
+    # get K R T
+    KRT_list = empty_param.get_KRT()
+    assert len(KRT_list) == 3
+    assert KRT_list[0].shape == (3, 3)
+    KRT_list = empty_param.get_KRT(k_dim=4)
+    assert KRT_list[0].shape == (4, 4)
+    with pytest.raises(ValueError):
+        KRT_list = empty_param.get_KRT(k_dim=5)