content: from typing import List, Optional, Tuple
  
  import numpy as np
  import torch
  from torch.nn import functional as F
  
  
  def get_position_map_from_depth(depth, mask, intrinsics, extrinsics, image_wh=None):
      """Compute the position map from the depth map and the camera parameters for a batch of views.
  
      Args:
          depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
          mask (torch.Tensor): The masks with the shape (B, H, W, 1).
          intrinsics (torch.Tensor): The camera intrinsics matrices with the shape (B, 3, 3).
          extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
          image_wh (Tuple[int, int]): The image width and height.
  
      Returns:
          torch.Tensor: The position maps with the shape (B, H, W, 3).
      """
      if image_wh is None:
          image_wh = depth.shape[2], depth.shape[1]
  
      B, H, W, _ = depth.shape
      depth = depth.squeeze(-1)
  
      u_coord, v_coord = torch.meshgrid(
          torch.arange(image_wh[0]), torch.arange(image_wh[1]), indexing="xy"
      )
      u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
      v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
  
      # Compute the position map by back-projecting depth pixels to 3D space
      x = (
          (u_coord - intrinsics[:, 0, 2].unsqueeze(-1).unsqueeze(-1))
          * depth
          / intrinsics[:, 0, 0].unsqueeze(-1).unsqueeze(-1)
      )
      y = (
          (v_coord - intrinsics[:, 1, 2].unsqueeze(-1).unsqueeze(-1))
          * depth
          / intrinsics[:, 1, 1].unsqueeze(-1).unsqueeze(-1)
      )
      z = depth
  
      # Concatenate to form the 3D coordinates in the camera frame
      camera_coords = torch.stack([x, y, z], dim=-1)
  
      # Apply the extrinsic matrix to get coordinates in the world frame
      coords_homogeneous = torch.nn.functional.pad(
          camera_coords, (0, 1), "constant", 1.0
      )  # Add a homogeneous coordinate
      world_coords = torch.matmul(
          coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
      ).view(B, H, W, 4)
  
      # Apply the mask to the position map
      position_map = world_coords[..., :3] * mask
  
      return position_map
  
  
  def get_position_map_from_depth_ortho(
      depth, mask, extrinsics, ortho_scale, image_wh=None
  ):
      """Compute the position map from the depth map and the camera parameters for a batch of views
      using orthographic projection with a given ortho_scale.
  
      Args:
          depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
          mask (torch.Tensor): The masks with the shape (B, H, W, 1).
          extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
          ortho_scale (torch.Tensor): The scaling factor for the orthographic projection with the shape (B, 1, 1, 1).
          image_wh (Tuple[int, int]): Optional. The image width and height.
  
      Returns:
          torch.Tensor: The position maps with the shape (B, H, W, 3).
      """
      if image_wh is None:
          image_wh = depth.shape[2], depth.shape[1]
  
      B, H, W, _ = depth.shape
      depth = depth.squeeze(-1)
  
      # Generating grid of coordinates in the image space
      u_coord, v_coord = torch.meshgrid(
          torch.arange(0, image_wh[0]), torch.arange(0, image_wh[1]), indexing="xy"
      )
      u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
      v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
  
      # Compute the position map using orthographic projection with ortho_scale
      x = (u_coord - image_wh[0] / 2) * ortho_scale / image_wh[0]
      y = (v_coord - image_wh[1] / 2) * ortho_scale / image_wh[1]
      z = depth
  
      # Concatenate to form the 3D coordinates in the camera frame
      camera_coords = torch.stack([x, y, z], dim=-1)
  
      # Apply the extrinsic matrix to get coordinates in the world frame
      coords_homogeneous = torch.nn.functional.pad(
          camera_coords, (0, 1), "constant", 1.0
      )  # Add a homogeneous coordinate
      world_coords = torch.matmul(
          coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
      ).view(B, H, W, 4)
  
      # Apply the mask to the position map
      position_map = world_coords[..., :3] * mask
  
      return position_map
  
  
  def get_opencv_from_blender(matrix_world, fov=None, image_size=None):
      # convert matrix_world to opencv format extrinsics
      opencv_world_to_cam = matrix_world.inverse()
      opencv_world_to_cam[1, :] *= -1
      opencv_world_to_cam[2, :] *= -1
      R, T = opencv_world_to_cam[:3, :3], opencv_world_to_cam[:3, 3]
  
      if fov is None:  # orthographic camera
          return R, T
  
      R, T = R.unsqueeze(0), T.unsqueeze(0)
      # convert fov to opencv format intrinsics
      focal = 1 / np.tan(fov / 2)
      intrinsics = np.diag(np.array([focal, focal, 1])).astype(np.float32)
      opencv_cam_matrix = (
          torch.from_numpy(intrinsics).unsqueeze(0).float().to(matrix_world.device)
      )
      opencv_cam_matrix[:, :2, -1] += torch.tensor([image_size / 2, image_size / 2]).to(
          matrix_world.device
      )
      opencv_cam_matrix[:, [0, 1], [0, 1]] *= image_size / 2
  
      return R, T, opencv_cam_matrix
  
  
  def get_ray_directions(
      H: int,
      W: int,
      focal: float,
      principal: Optional[Tuple[float, float]] = None,
      use_pixel_centers: bool = True,
  ) -> torch.Tensor:
      """
      Get ray directions for all pixels in camera coordinate.
      Args:
          H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
      Outputs:
          directions: (H, W, 3), the direction of the rays in camera coordinate
      """
      pixel_center = 0.5 if use_pixel_centers else 0
      cx, cy = W / 2, H / 2 if principal is None else principal
      i, j = torch.meshgrid(
          torch.arange(W, dtype=torch.float32) + pixel_center,
          torch.arange(H, dtype=torch.float32) + pixel_center,
          indexing="xy",
      )
      directions = torch.stack(
          [(i - cx) / focal, -(j - cy) / focal, -torch.ones_like(i)], -1
      )
      return F.normalize(directions, dim=-1)
  
  
  def get_rays(
      directions: torch.Tensor, c2w: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Get ray origins and directions from camera coordinates to world coordinates
      Args:
          directions: (H, W, 3) ray directions in camera coordinates
          c2w: (4, 4) camera-to-world transformation matrix
      Outputs:
          rays_o, rays_d: (H, W, 3) ray origins and directions in world coordinates
      """
      # Rotate ray directions from camera coordinate to the world coordinate
      rays_d = directions @ c2w[:3, :3].T
      rays_o = c2w[:3, 3].expand(rays_d.shape)
      return rays_o, rays_d
  
  
  def compute_plucker_embed(
      c2w: torch.Tensor, image_width: int, image_height: int, focal: float
  ) -> torch.Tensor:
      """
      Computes Plucker coordinates for a camera.
      Args:
          c2w: (4, 4) camera-to-world transformation matrix
          image_width: Image width
          image_height: Image height
          focal: Focal length of the camera
      Returns:
          plucker: (6, H, W) Plucker embedding
      """
      directions = get_ray_directions(image_height, image_width, focal)
      rays_o, rays_d = get_rays(directions, c2w)
      # Cross product to get Plucker coordinates
      cross = torch.cross(rays_o, rays_d, dim=-1)
      plucker = torch.cat((rays_d, cross), dim=-1)
      return plucker.permute(2, 0, 1)
  
  
  def get_plucker_embeds_from_cameras(
      c2w: List[torch.Tensor], fov: List[float], image_size: int
  ) -> torch.Tensor:
      """
      Given lists of camera transformations and fov, returns the batched plucker embeddings.
      Args:
          c2w: list of camera-to-world transformation matrices
          fov: list of field of view values
          image_size: size of the image
      Returns:
          plucker_embeds: (B, 6, H, W) batched plucker embeddings
      """
      plucker_embeds = []
      for cam_matrix, cam_fov in zip(c2w, fov):
          focal = 0.5 * image_size / np.tan(0.5 * cam_fov)
          plucker = compute_plucker_embed(cam_matrix, image_size, image_size, focal)
          plucker_embeds.append(plucker)
      return torch.stack(plucker_embeds)
  
  
  def get_plucker_embeds_from_cameras_ortho(
      c2w: List[torch.Tensor], ortho_scale: List[float], image_size: int
  ):
      """
      Given lists of camera transformations and fov, returns the batched plucker embeddings.
  
      Parameters:
          c2w: list of camera-to-world transformation matrices
          fov: list of field of view values
          image_size: size of the image
  
      Returns:
          plucker_embeds: plucker embeddings (B, 6, H, W)
      """
      plucker_embeds = []
      # compute pairwise mask and plucker embeddings
      for cam_matrix, scale in zip(c2w, ortho_scale):
          # blender to opencv to pytorch3d
          R, T = get_opencv_from_blender(cam_matrix)
          cam_pos = -R.T @ T
          view_dir = R.T @ torch.tensor([0, 0, 1]).float().to(cam_matrix.device)
          # normalize camera position
          cam_pos = F.normalize(cam_pos, dim=0)
          plucker = torch.concat([view_dir, cam_pos])
          plucker = plucker.unsqueeze(-1).unsqueeze(-1).repeat(1, image_size, image_size)
          plucker_embeds.append(plucker)
  
      plucker_embeds = torch.stack(plucker_embeds)
  
      return plucker_embeds
  