text: # Copyright (c) Facebook, Inc. and its affiliates.
  # All rights reserved.
  #
  # This source code is licensed under the BSD-style license found in the
  # LICENSE file in the root directory of this source tree.
  
  import unittest
  
  import numpy as np
  import torch
  from common_testing import TestCaseMixin
  from pytorch3d.ops import corresponding_cameras_alignment
  from pytorch3d.renderer.cameras import (
      OpenGLOrthographicCameras,
      OpenGLPerspectiveCameras,
      SfMOrthographicCameras,
      SfMPerspectiveCameras,
  )
  from pytorch3d.transforms.rotation_conversions import random_rotations
  from pytorch3d.transforms.so3 import so3_exp_map, so3_relative_angle
  from test_cameras import init_random_cameras
  
  
  class TestCamerasAlignment(TestCaseMixin, unittest.TestCase):
      def setUp(self) -> None:
          super().setUp()
          torch.manual_seed(42)
          np.random.seed(42)
  
      def test_corresponding_cameras_alignment(self):
          """
          Checks the corresponding_cameras_alignment function.
          """
          self.skipTest("Temporarily disabled pending investigation")
          device = torch.device("cuda:0")
  
          # try few different random setups
          for _ in range(3):
              for estimate_scale in (True, False):
                  # init true alignment transform
                  R_align_gt = random_rotations(1, device=device)[0]
                  T_align_gt = torch.randn(3, dtype=torch.float32, device=device)
  
                  # init true scale
                  if estimate_scale:
                      s_align_gt = torch.randn(
                          1, dtype=torch.float32, device=device
                      ).exp()
                  else:
                      s_align_gt = torch.tensor(1.0, dtype=torch.float32, device=device)
  
                  for cam_type in (
                      SfMOrthographicCameras,
                      OpenGLPerspectiveCameras,
                      OpenGLOrthographicCameras,
                      SfMPerspectiveCameras,
                  ):
                      # try well-determined and underdetermined cases
                      for batch_size in (10, 4, 3, 2, 1):
                          # get random cameras
                          cameras = init_random_cameras(
                              cam_type, batch_size, random_z=True
                          ).to(device)
                          # try all alignment modes
                          for mode in ("extrinsics", "centers"):
                              # try different noise levels
                              for add_noise in (0.0, 0.01, 1e-4):
                                  self._corresponding_cameras_alignment_test_case(
                                      cameras,
                                      R_align_gt,
                                      T_align_gt,
                                      s_align_gt,
                                      estimate_scale,
                                      mode,
                                      add_noise,
                                  )
  
      def _corresponding_cameras_alignment_test_case(
          self,
          cameras,
          R_align_gt,
          T_align_gt,
          s_align_gt,
          estimate_scale,
          mode,
          add_noise,
      ):
          batch_size = cameras.R.shape[0]
  
          # get target camera centers
          R_new = torch.bmm(R_align_gt[None].expand_as(cameras.R), cameras.R)
          T_new = (
              torch.bmm(T_align_gt[None, None].repeat(batch_size, 1, 1), cameras.R)[:, 0]
              + cameras.T
          ) * s_align_gt
  
          if add_noise != 0.0:
              R_new = torch.bmm(R_new, so3_exp_map(torch.randn_like(T_new) * add_noise))
              T_new += torch.randn_like(T_new) * add_noise
  
          # create new cameras from R_new and T_new
          cameras_tgt = cameras.clone()
          cameras_tgt.R = R_new
          cameras_tgt.T = T_new
  
          # align cameras and cameras_tgt
          cameras_aligned = corresponding_cameras_alignment(
              cameras, cameras_tgt, estimate_scale=estimate_scale, mode=mode
          )
  
          if batch_size <= 2 and mode == "centers":
              # underdetermined case - check only the center alignment error
              # since the rotation and translation are ambiguous here
              self.assertClose(
                  cameras_aligned.get_camera_center(),
                  cameras_tgt.get_camera_center(),
                  atol=max(add_noise * 7.0, 1e-4),
              )
  
          else:
  
              def _rmse(a):
                  return (torch.norm(a, dim=1, p=2) ** 2).mean().sqrt()
  
              if add_noise != 0.0:
                  # in a noisy case check mean rotation/translation error for
                  # extrinsic alignment and root mean center error for center alignment
                  if mode == "centers":
                      self.assertNormsClose(
                          cameras_aligned.get_camera_center(),
                          cameras_tgt.get_camera_center(),
                          _rmse,
                          atol=max(add_noise * 10.0, 1e-4),
                      )
                  elif mode == "extrinsics":
                      angle_err = so3_relative_angle(
                          cameras_aligned.R, cameras_tgt.R
                      ).mean()
                      self.assertClose(
                          angle_err, torch.zeros_like(angle_err), atol=add_noise * 10.0
                      )
                      self.assertNormsClose(
                          cameras_aligned.T, cameras_tgt.T, _rmse, atol=add_noise * 7.0
                      )
                  else:
                      raise ValueError(mode)
  
              else:
                  # compare the rotations and translations of cameras
                  self.assertClose(cameras_aligned.R, cameras_tgt.R, atol=3e-4)
                  self.assertClose(cameras_aligned.T, cameras_tgt.T, atol=3e-4)
                  # compare the centers
                  self.assertClose(
                      cameras_aligned.get_camera_center(),
                      cameras_tgt.get_camera_center(),
                      atol=3e-4,
                  )
  
      @staticmethod
      def corresponding_cameras_alignment(
          batch_size: int, estimate_scale: bool, mode: str, cam_type=SfMPerspectiveCameras
      ):
          device = torch.device("cuda:0")
          cameras_src, cameras_tgt = [
              init_random_cameras(cam_type, batch_size, random_z=True).to(device)
              for _ in range(2)
          ]
  
          torch.cuda.synchronize()
  
          def compute_corresponding_cameras_alignment():
              corresponding_cameras_alignment(
                  cameras_src, cameras_tgt, estimate_scale=estimate_scale, mode=mode
              )
              torch.cuda.synchronize()
  
          return compute_corresponding_cameras_alignment
  