text: # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
  
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  
  # pyre-fixme[21]: Could not find name `_C` in `pytorch3d`.
  from pytorch3d import _C
  from pytorch3d.renderer.mesh.rasterize_meshes import pix_to_non_square_ndc
  
  
  # Maxinum number of faces per bins for
  # coarse-to-fine rasterization
  kMaxPointsPerBin = 22
  
  
  def rasterize_points(
      pointclouds,
      image_size: Union[int, List[int], Tuple[int, int]] = 256,
      radius: Union[float, List, Tuple, torch.Tensor] = 0.01,
      points_per_pixel: int = 8,
      bin_size: Optional[int] = None,
      max_points_per_bin: Optional[int] = None,
  ):
      """
      Each pointcloud is rasterized onto a separate image of shape
      (H, W) if `image_size` is a tuple or (image_size, image_size) if it
      is an int.
  
      If the desired image size is non square (i.e. a tuple of (H, W) where H != W)
      the aspect ratio needs special consideration. There are two aspect ratios
      to be aware of:
          - the aspect ratio of each pixel
          - the aspect ratio of the output image
      The camera can be used to set the pixel aspect ratio. In the rasterizer,
      we assume square pixels, but variable image aspect ratio (i.e rectangle images).
  
      In most cases you will want to set the camera aspect ratio to
      1.0 (i.e. square pixels) and only vary the
      `image_size` (i.e. the output image dimensions in pix
  
      Args:
          pointclouds: A Pointclouds object representing a batch of point clouds to be
              rasterized. This is a batch of N pointclouds, where each point cloud
              can have a different number of points; the coordinates of each point
              are (x, y, z). The coordinates are expected to
              be in normalized device coordinates (NDC): [-1, 1]^3 with the camera at
              (0, 0, 0); In the camera coordinate frame the x-axis goes from right-to-left,
              the y-axis goes from bottom-to-top, and the z-axis goes from back-to-front.
          image_size: Size in pixels of the output image to be rasterized.
              Can optionally be a tuple of (H, W) in the case of non square images.
          radius (Optional): The radius (in NDC units) of the disk to
              be rasterized. This can either be a float in which case the same radius is used
              for each point, or a torch.Tensor of shape (N, P) giving a radius per point
              in the batch.
          points_per_pixel (Optional): We will keep track of this many points per
              pixel, returning the nearest points_per_pixel points along the z-axis
          bin_size: Size of bins to use for coarse-to-fine rasterization. Setting
              bin_size=0 uses naive rasterization; setting bin_size=None attempts to
              set it heuristically based on the shape of the input. This should not
              affect the output, but can affect the speed of the forward pass.
          points_per_bin: Only applicable when using coarse-to-fine rasterization
              (bin_size > 0); this is the maxiumum number of points allowed within each
              bin. If more than this many points actually fall into a bin, an error
              will be raised. This should not affect the output values, but can affect
              the memory usage in the forward pass.
  
      Returns:
          3-element tuple containing
  
          - **idx**: int32 Tensor of shape (N, image_size, image_size, points_per_pixel)
            giving the indices of the nearest points at each pixel, in ascending
            z-order. Concretely `idx[n, y, x, k] = p` means that `points[p]` is the kth
            closest point (along the z-direction) to pixel (y, x) - note that points
            represents the packed points of shape (P, 3).
            Pixels that are hit by fewer than points_per_pixel are padded with -1.
          - **zbuf**: Tensor of shape (N, image_size, image_size, points_per_pixel)
            giving the z-coordinates of the nearest points at each pixel, sorted in
            z-order. Concretely, if `idx[n, y, x, k] = p` then
            `zbuf[n, y, x, k] = points[n, p, 2]`. Pixels hit by fewer than
            points_per_pixel are padded with -1
          - **dists2**: Tensor of shape (N, image_size, image_size, points_per_pixel)
            giving the squared Euclidean distance (in NDC units) in the x/y plane
            for each point closest to the pixel. Concretely if `idx[n, y, x, k] = p`
            then `dists[n, y, x, k]` is the squared distance between the pixel (y, x)
            and the point `(points[n, p, 0], points[n, p, 1])`. Pixels hit with fewer
            than points_per_pixel are padded with -1.
  
          In the case that image_size is a tuple of (H, W) then the outputs
          will be of shape `(N, H, W, ...)`.
      """
      points_packed = pointclouds.points_packed()
      cloud_to_packed_first_idx = pointclouds.cloud_to_packed_first_idx()
      num_points_per_cloud = pointclouds.num_points_per_cloud()
  
      radius = _format_radius(radius, pointclouds)
  
      # In the case that H != W use the max image size to set the bin_size
      # to accommodate the num bins constraint in the coarse rasteizer.
      # If the ratio of H:W is large this might cause issues as the smaller
      # dimension will have fewer bins.
      # TODO: consider a better way of setting the bin size.
      if isinstance(image_size, (tuple, list)):
          if len(image_size) != 2:
              raise ValueError("Image size can only be a tuple/list of (H, W)")
          if not all(i > 0 for i in image_size):
              raise ValueError(
                  "Image sizes must be greater than 0; got %d, %d" % image_size
              )
          if not all(type(i) == int for i in image_size):
              raise ValueError("Image sizes must be integers; got %f, %f" % image_size)
          max_image_size = max(*image_size)
          im_size = image_size
      else:
          im_size = (image_size, image_size)
          max_image_size = image_size
  
      if bin_size is None:
          if not points_packed.is_cuda:
              # Binned CPU rasterization not fully implemented
              bin_size = 0
          else:
              bin_size = int(2 ** max(np.ceil(np.log2(max_image_size)) - 4, 4))
  
      if bin_size != 0:
          # There is a limit on the number of points per bin in the cuda kernel.
          points_per_bin = 1 + (max_image_size - 1) // bin_size
          if points_per_bin >= kMaxPointsPerBin:
              raise ValueError(
                  "bin_size too small, number of points per bin must be less than %d; got %d"
                  % (kMaxPointsPerBin, points_per_bin)
              )
  
      if max_points_per_bin is None:
          max_points_per_bin = int(max(10000, pointclouds._P / 5))
  
      # Function.apply cannot take keyword args, so we handle defaults in this
      # wrapper and call apply with positional args only
      # pyre-fixme[16]: `_RasterizePoints` has no attribute `apply`.
      return _RasterizePoints.apply(
          points_packed,
          cloud_to_packed_first_idx,
          num_points_per_cloud,
          im_size,
          radius,
          points_per_pixel,
          bin_size,
          max_points_per_bin,
      )
  
  
  def _format_radius(
      radius: Union[float, List, Tuple, torch.Tensor], pointclouds
  ) -> torch.Tensor:
      """
      Format the radius as a torch tensor of shape (P_packed,)
      where P_packed is the total number of points in the
      batch (i.e. pointclouds.points_packed().shape[0]).
  
      This will enable support for a different size radius
      for each point in the batch.
  
      Args:
          radius: can be a float, List, Tuple or tensor of
              shape (N, P_padded) where P_padded is the
              maximum number of points for each pointcloud
              in the batch.
  
      Returns:
          radius: torch.Tensor of shape (P_packed)
      """
      N, P_padded = pointclouds._N, pointclouds._P
      points_packed = pointclouds.points_packed()
      P_packed = points_packed.shape[0]
      if isinstance(radius, (list, tuple)):
          radius = torch.tensor(radius).type_as(points_packed)
      if isinstance(radius, torch.Tensor):
          if N == 1 and radius.ndim == 1:
              radius = radius[None, ...]
          if radius.shape != (N, P_padded):
              msg = "radius must be of shape (N, P): got %s"
              raise ValueError(msg % (repr(radius.shape)))
          else:
              padded_to_packed_idx = pointclouds.padded_to_packed_idx()
              radius = radius.view(-1)[padded_to_packed_idx]
      elif isinstance(radius, float):
          radius = torch.full((P_packed,), fill_value=radius).type_as(points_packed)
      else:
          msg = "radius must be a float, list, tuple or tensor; got %s"
          raise ValueError(msg % type(radius))
      return radius
  
  
  class _RasterizePoints(torch.autograd.Function):
      @staticmethod
      # pyre-fixme[14]: `forward` overrides method defined in `Function` inconsistently.
      # pyre-fixme[14]: `forward` overrides method defined in `Function` inconsistently.
      def forward(
          ctx,
          points,  # (P, 3)
          cloud_to_packed_first_idx,
          num_points_per_cloud,
          image_size: Union[List[int], Tuple[int, int]] = (256, 256),
          radius: Union[float, torch.Tensor] = 0.01,
          points_per_pixel: int = 8,
          bin_size: int = 0,
          max_points_per_bin: int = 0,
      ):
          # TODO: Add better error handling for when there are more than
          # max_points_per_bin in any bin.
          args = (
              points,
              cloud_to_packed_first_idx,
              num_points_per_cloud,
              image_size,
              radius,
              points_per_pixel,
              bin_size,
              max_points_per_bin,
          )
          # pyre-fixme[16]: Module `pytorch3d` has no attribute `_C`.
          idx, zbuf, dists = _C.rasterize_points(*args)
          ctx.save_for_backward(points, idx)
          ctx.mark_non_differentiable(idx)
          return idx, zbuf, dists
  
      @staticmethod
      def backward(ctx, grad_idx, grad_zbuf, grad_dists):
          grad_points = None
          grad_cloud_to_packed_first_idx = None
          grad_num_points_per_cloud = None
          grad_image_size = None
          grad_radius = None
          grad_points_per_pixel = None
          grad_bin_size = None
          grad_max_points_per_bin = None
          points, idx = ctx.saved_tensors
          args = (points, idx, grad_zbuf, grad_dists)
          grad_points = _C.rasterize_points_backward(*args)
          grads = (
              grad_points,
              grad_cloud_to_packed_first_idx,
              grad_num_points_per_cloud,
              grad_image_size,
              grad_radius,
              grad_points_per_pixel,
              grad_bin_size,
              grad_max_points_per_bin,
          )
          return grads
  
  
  def rasterize_points_python(
      pointclouds,
      image_size: Union[int, Tuple[int, int]] = 256,
      radius: Union[float, torch.Tensor] = 0.01,
      points_per_pixel: int = 8,
  ):
      """
      Naive pure PyTorch implementation of pointcloud rasterization.
  
      Inputs / Outputs: Same as above
      """
      N = len(pointclouds)
      H, W = (
          image_size
          if isinstance(image_size, (tuple, list))
          else (image_size, image_size)
      )
      K = points_per_pixel
      device = pointclouds.device
  
      points_packed = pointclouds.points_packed()
      cloud_to_packed_first_idx = pointclouds.cloud_to_packed_first_idx()
      num_points_per_cloud = pointclouds.num_points_per_cloud()
  
      # Support variable size radius for each point in the batch
      radius = _format_radius(radius, pointclouds)
  
      # Intialize output tensors.
      point_idxs = torch.full(
          (N, H, W, K), fill_value=-1, dtype=torch.int32, device=device
      )
      zbuf = torch.full((N, H, W, K), fill_value=-1, dtype=torch.float32, device=device)
      pix_dists = torch.full(
          (N, H, W, K), fill_value=-1, dtype=torch.float32, device=device
      )
  
      # NDC is from [-1, 1]. Get pixel size using specified image size.
      radius2 = radius * radius
  
      # Iterate through the batch of point clouds.
      for n in range(N):
          point_start_idx = cloud_to_packed_first_idx[n]
          point_stop_idx = point_start_idx + num_points_per_cloud[n]
  
          # Iterate through the horizontal lines of the image from top to bottom.
          for yi in range(H):
              # Y coordinate of one end of the image. Reverse the ordering
              # of yi so that +Y is pointing up in the image.
              yfix = H - 1 - yi
              yf = pix_to_non_square_ndc(yfix, H, W)
  
              # Iterate through pixels on this horizontal line, left to right.
              for xi in range(W):
                  # X coordinate of one end of the image. Reverse the ordering
                  # of xi so that +X is pointing to the left in the image.
                  xfix = W - 1 - xi
                  xf = pix_to_non_square_ndc(xfix, W, H)
  
                  top_k_points = []
                  # Check whether each point in the batch affects this pixel.
                  for p in range(point_start_idx, point_stop_idx):
                      px, py, pz = points_packed[p, :]
                      r = radius2[p]
                      if pz < 0:
                          continue
                      dx = px - xf
                      dy = py - yf
                      dist2 = dx * dx + dy * dy
                      if dist2 < r:
                          top_k_points.append((pz, p, dist2))
                          top_k_points.sort()
                          if len(top_k_points) > K:
                              top_k_points = top_k_points[:K]
                  for k, (pz, p, dist2) in enumerate(top_k_points):
                      zbuf[n, yi, xi, k] = pz
                      point_idxs[n, yi, xi, k] = p
                      pix_dists[n, yi, xi, k] = dist2
      return point_idxs, zbuf, pix_dists
  