models/rendering/point_sampler.py

# python3.7
"""Contains the functions to sample points in 3D space."""

import numpy as np

import torch
import torch.nn.functional as F

__all__ = [
    'PointSampler'
]

_POINT_SAMPLING_STRATEGIES = [
    'uniform', 'normal', 'ray_dependent', 'point_dependent'
]

_POINT_PERTURBING_STRATEGIES = [
    'no', 'middle_uniform', 'uniform', 'self_uniform'
]

_TENSOR_SAMPLING_STRATEGIES = [
    'fix', 'uniform', 'normal', 'hybrid', 'truncated_normal'
]


class PointSampler(torch.nn.Module):
    """Defines the class to help sample points.

    This class implements the `forward()` function for point sampling, which
    includes the following steps:

    1. Sample rays in the camera coordinate system.
    2. Sample points on each ray.
    3. Perturb points on each ray.
    4. Sample camera extrinsics.
    5. Transform points to the world coordinate system.
    """

    def __init__(self,
                 num_points=16,
                 fov=30,
                 image_boundary_value=1.0,
                 cam_look_at_dir=-1,
                 pixel_center=False,
                 y_descending=True,
                 # Point sampling (i.e., radial distance w.r.t. camera) related.
                 sampling_strategy='uniform',
                 focal=None,
                 dis_min=None,
                 dis_max=None,
                 dis_mean=None,
                 dis_stddev=None,
                 per_ray_ref=None,
                 per_point_ref=None,
                 perturbation_strategy='middle_uniform',
                 # Camera sampling related.
                 radius_strategy='fix',
                 radius_fix=None,
                 radius_min=None,
                 radius_max=None,
                 radius_mean=None,
                 radius_stddev=None,
                 polar_strategy='uniform',
                 polar_fix=None,
                 polar_min=None,
                 polar_max=None,
                 polar_mean=None,
                 polar_stddev=None,
                 azimuthal_strategy='uniform',
                 azimuthal_fix=None,
                 azimuthal_min=None,
                 azimuthal_max=None,
                 azimuthal_mean=None,
                 azimuthal_stddev=None,
                 use_spherical_uniform_position=False,
                 pitch_strategy='fix',
                 pitch_fix=0,
                 pitch_min=None,
                 pitch_max=None,
                 pitch_mean=None,
                 pitch_stddev=None,
                 yaw_strategy='fix',
                 yaw_fix=0,
                 yaw_min=None,
                 yaw_max=None,
                 yaw_mean=None,
                 yaw_stddev=None,
                 roll_strategy='fix',
                 roll_fix=0,
                 roll_min=None,
                 roll_max=None,
                 roll_mean=None,
                 roll_stddev=None):
        """Initializes hyper-parameters for point sampling.

        Detailed description of each argument can be found in functions
        `get_ray_per_pixel()`, `sample_points_per_ray()`,
        `perturb_points_per_ray()`, `sample_camera_extrinsics()`.
        """
        super().__init__()
        self.num_points = num_points
        self.fov = fov
        self.image_boundary_value = image_boundary_value
        self.cam_look_at_dir = cam_look_at_dir
        self.pixel_center = pixel_center
        self.y_descending = y_descending

        self.sampling_strategy = sampling_strategy
        self.dis_min = dis_min
        self.dis_max = dis_max
        self.dis_mean = dis_mean
        self.dis_stddev = dis_stddev
        self.per_ray_ref = per_ray_ref
        self.per_point_ref = per_point_ref
        self.perturbation_strategy = perturbation_strategy

        self.radius_strategy = radius_strategy
        self.radius_fix = radius_fix
        self.radius_min = radius_min
        self.radius_max = radius_max
        self.radius_mean = radius_mean
        self.radius_stddev = radius_stddev
        self.polar_strategy = polar_strategy
        self.polar_fix = polar_fix
        self.polar_min = polar_min
        self.polar_max = polar_max
        self.polar_mean = polar_mean
        self.polar_stddev = polar_stddev
        self.azimuthal_strategy = azimuthal_strategy
        self.azimuthal_fix = azimuthal_fix
        self.azimuthal_min = azimuthal_min
        self.azimuthal_max = azimuthal_max
        self.azimuthal_mean = azimuthal_mean
        self.azimuthal_stddev = azimuthal_stddev
        self.use_spherical_uniform_position = use_spherical_uniform_position
        self.pitch_strategy = pitch_strategy
        self.pitch_fix = pitch_fix
        self.pitch_min = pitch_min
        self.pitch_max = pitch_max
        self.pitch_mean = pitch_mean
        self.pitch_stddev = pitch_stddev
        self.yaw_strategy = yaw_strategy
        self.yaw_fix = yaw_fix
        self.yaw_min = yaw_min
        self.yaw_max = yaw_max
        self.yaw_mean = yaw_mean
        self.yaw_stddev = yaw_stddev
        self.roll_strategy = roll_strategy
        self.roll_fix = roll_fix
        self.roll_min = roll_min
        self.roll_max = roll_max
        self.roll_mean = roll_mean
        self.roll_stddev = roll_stddev
        self.focal = focal

    def forward(self,
                batch_size,
                image_size,
                focal=None,
                cam2world_matrix=None,
                **kwargs):
        """Samples points.

        `K` denotes the number of points on each ray.

        Args:
            batch_size: Batch size of images. Denoted as `N`.
            image_size: Size of the image. One element indicates square image,
                while two elements stand for height and width respectively.
                Denoted as `H` and `W`.
            **kwargs: Additional keyword arguments to override the variables
                initialized in `__init__()`.

        Returns:
            A dictionary, containing
                - `camera_radius`: camera radius w.r.t. the world coordinate
                    system, with shape [N].
                - `camera_polar`: camera polar w.r.t. the world coordinate
                    system, with shape [N].
                - `camera_azimuthal`: camera azimuthal w.r.t. the world
                    coordinate system, with shape [N].
                - `camera_pitch`: camera pitch w.r.t. the camera coordinate
                    system, with shape [N].
                - `camera_yaw`: camera yaw w.r.t. the camera coordinate system,
                    with shape [N].
                - `camera_roll`: camera roll w.r.t. the camera coordinate
                    system, with shape [N].
                - `camera_pos`: camera position, i.e., the (x, y, z) coordinate
                    in the world coordinate system, with shape [N, 3].
                - `cam2world_matrix`: transformation matrix to transform the
                    camera coordinate system to the world coordinate system,
                    with shape [N, 4, 4].
                - `rays_camera`: ray directions in the camera coordinate system,
                    with shape [N, H, W, 3].
                - `rays_world`: ray directions in the world coordinate system,
                    with shape [N, H, W, 3].
                - `radii_raw`: raw per-point radial distance w.r.t. the camera
                    position, with shape [N, H, W, K].
                - `radii`: per-point radial distance after perturbation w.r.t.
                    the camera position, with shape [N, H, W, K].
                - `points_camera`: per-point coordinate in the camera coordinate
                    system, with shape [N, H, W, K, 3].
                - `points_world`: per-point coordinate in the world coordinate
                    system, with shape [N, H, W, K, 3].
        """
        num_points = kwargs.get('num_points', self.num_points)
        fov = kwargs.get('fov', self.fov)
        image_boundary_value = kwargs.get(
            'image_boundary_value', self.image_boundary_value)
        cam_look_at_dir = kwargs.get('cam_look_at_dir', self.cam_look_at_dir)
        pixel_center = kwargs.get('pixel_center', self.pixel_center)
        y_descending = kwargs.get('y_descending', self.y_descending)
        sampling_strategy = kwargs.get(
            'sampling_strategy', self.sampling_strategy)
        dis_min = kwargs.get('dis_min', self.dis_min)
        dis_max = kwargs.get('dis_max', self.dis_max)
        dis_mean = kwargs.get('dis_mean', self.dis_mean)
        dis_stddev = kwargs.get('dis_stddev', self.dis_stddev)
        per_ray_ref = kwargs.get('per_ray_ref', self.per_ray_ref)
        per_point_ref = kwargs.get('per_point_ref', self.per_point_ref)
        perturbation_strategy = kwargs.get(
            'perturbation_strategy', self.perturbation_strategy)
        radius_strategy = kwargs.get('radius_strategy', self.radius_strategy)
        radius_fix = kwargs.get('radius_fix', self.radius_fix)
        radius_min = kwargs.get('radius_min', self.radius_min)
        radius_max = kwargs.get('radius_max', self.radius_max)
        radius_mean = kwargs.get('radius_mean', self.radius_mean)
        radius_stddev = kwargs.get('radius_stddev', self.radius_stddev)
        polar_strategy = kwargs.get('polar_strategy', self.polar_strategy)
        polar_fix = kwargs.get('polar_fix', self.polar_fix)
        polar_min = kwargs.get('polar_min', self.polar_min)
        polar_max = kwargs.get('polar_max', self.polar_max)
        polar_mean = kwargs.get('polar_mean', self.polar_mean)
        polar_stddev = kwargs.get('polar_stddev', self.polar_stddev)
        azimuthal_strategy = kwargs.get(
            'azimuthal_strategy', self.azimuthal_strategy)
        azimuthal_fix = kwargs.get('azimuthal_fix', self.azimuthal_fix)
        azimuthal_min = kwargs.get('azimuthal_min', self.azimuthal_min)
        azimuthal_max = kwargs.get('azimuthal_max', self.azimuthal_max)
        azimuthal_mean = kwargs.get('azimuthal_mean', self.azimuthal_mean)
        azimuthal_stddev = kwargs.get('azimuthal_stddev', self.azimuthal_stddev)
        use_spherical_uniform_position = kwargs.get(
            'use_spherical_uniform_position',
            self.use_spherical_uniform_position)
        pitch_strategy = kwargs.get('pitch_strategy', self.pitch_strategy)
        pitch_fix = kwargs.get('pitch_fix', self.pitch_fix)
        pitch_min = kwargs.get('pitch_min', self.pitch_min)
        pitch_max = kwargs.get('pitch_max', self.pitch_max)
        pitch_mean = kwargs.get('pitch_mean', self.pitch_mean)
        pitch_stddev = kwargs.get('pitch_stddev', self.pitch_stddev)
        yaw_strategy = kwargs.get('yaw_strategy', self.yaw_strategy)
        yaw_fix = kwargs.get('yaw_fix', self.yaw_fix)
        yaw_min = kwargs.get('yaw_min', self.yaw_min)
        yaw_max = kwargs.get('yaw_max', self.yaw_max)
        yaw_mean = kwargs.get('yaw_mean', self.yaw_mean)
        yaw_stddev = kwargs.get('yaw_stddev', self.yaw_stddev)
        roll_strategy = kwargs.get('roll_strategy', self.roll_strategy)
        roll_fix = kwargs.get('roll_fix', self.roll_fix)
        roll_min = kwargs.get('roll_min', self.roll_min)
        roll_max = kwargs.get('roll_max', self.roll_max)
        roll_mean = kwargs.get('roll_mean', self.roll_mean)
        roll_stddev = kwargs.get('roll_stddev', self.roll_stddev)

        rays_camera = get_ray_per_pixel(batch_size=batch_size,
                                        image_size=image_size,
                                        fov=fov,
                                        boundary=image_boundary_value,
                                        focal=focal,
                                        cam_look_at_dir=cam_look_at_dir,
                                        pixel_center=pixel_center,
                                        y_descending=y_descending)

        radii_raw = sample_points_per_ray(batch_size=batch_size,
                                          image_size=image_size,
                                          num_points=num_points,
                                          strategy=sampling_strategy,
                                          dis_min=dis_min,
                                          dis_max=dis_max,
                                          dis_mean=dis_mean,
                                          dis_stddev=dis_stddev,
                                          per_ray_ref=per_ray_ref,
                                          per_point_ref=per_point_ref)
        radii = perturb_points_per_ray(radii=radii_raw,
                                       strategy=perturbation_strategy)

        camera_info = {}
        if cam2world_matrix is not None:
            camera_info.update(dict(
                cam2world_matrix=cam2world_matrix,
                radius=None,
                polar=None,
                azimuthal=None,
                pitch=None,
                yaw=None,
                roll=None,
                camera_pos=None,
            ))
        else:
            camera_info = sample_camera_extrinsics(
                batch_size=batch_size,
                radius_strategy=radius_strategy,
                radius_fix=radius_fix,
                radius_min=radius_min,
                radius_max=radius_max,
                radius_mean=radius_mean,
                radius_stddev=radius_stddev,
                polar_strategy=polar_strategy,
                polar_fix=polar_fix,
                polar_min=polar_min,
                polar_max=polar_max,
                polar_mean=polar_mean,
                polar_stddev=polar_stddev,
                azimuthal_strategy=azimuthal_strategy,
                azimuthal_fix=azimuthal_fix,
                azimuthal_min=azimuthal_min,
                azimuthal_max=azimuthal_max,
                azimuthal_mean=azimuthal_mean,
                azimuthal_stddev=azimuthal_stddev,
                use_spherical_uniform_position=use_spherical_uniform_position,
                pitch_strategy=pitch_strategy,
                pitch_fix=pitch_fix,
                pitch_min=pitch_min,
                pitch_max=pitch_max,
                pitch_mean=pitch_mean,
                pitch_stddev=pitch_stddev,
                yaw_strategy=yaw_strategy,
                yaw_fix=yaw_fix,
                yaw_min=yaw_min,
                yaw_max=yaw_max,
                yaw_mean=yaw_mean,
                yaw_stddev=yaw_stddev,
                roll_strategy=roll_strategy,
                roll_fix=roll_fix,
                roll_min=roll_min,
                roll_max=roll_max,
                roll_mean=roll_mean,
                roll_stddev=roll_stddev)

        points = get_point_coord(
            rays_camera=rays_camera,
            radii=radii,
            cam2world_matrix=camera_info['cam2world_matrix'])

        return {
            'camera_radius': camera_info['radius'],  # [N]
            'camera_polar': camera_info['polar'],  # [N]
            'camera_azimuthal': camera_info['azimuthal'],  # [N]
            'camera_pitch': camera_info['pitch'],  # [N]
            'camera_yaw': camera_info['yaw'],  # [N]
            'camera_roll': camera_info['roll'],  # [N]
            'camera_pos':camera_info['camera_pos'],  # [N, 3]
            'cam2world_matrix': camera_info['cam2world_matrix'],  # [N, 4, 4]
            'rays_camera': rays_camera,  # [N, H, W, 3]
            'rays_world': points['rays_world'],  # [N, H, W, 3]
            'ray_origins_world': points['ray_origins_world'], # [N, H, W, 3]
            'radii_raw': radii_raw,  # [N, H, W, K]
            'radii': radii,  # [N, H, W, K]
            'points_camera': points['points_camera'],  # [N, H, W, K, 3]
            'points_world': points['points_world']  # [N, H, W, K, 3]
        }


def get_ray_per_pixel(batch_size,
                      image_size,
                      fov,
                      boundary=1.0,
                      focal=None,
                      cam_look_at_dir=-1,
                      pixel_center=False,
                      y_descending=True):
    """Gets ray direction for each image pixel under camera coordinate system.

    Each ray direction is represent by a vector, [x, y, z], under the following
    coordinate system:

    - The origin is set at the camera position.
    - The X axis is set as the horizontal direction of the image plane, with
      larger value on the right.
    - The Y axis is set as the vertical direction of the image plane, with
      larger value on the top.
    - The Z axis is set as the direction perpendicular to the image plane,
      from the image center pointing to the camera. In other words, the z
      coordinate of the image plane is negative.
    - The above coordinate system is a right-hand one.

    Taking a 5x5 image (with boundary 1.0) as an instance, the per-pixel (x, y)
    coordinates should look like:

    (-1.0,  1.0) (-0.5,  1.0) (0.0,  1.0) (0.5,  1.0) (1.0,  1.0)
    (-1.0,  0.5) (-0.5,  0.5) (0.0,  0.5) (0.5,  0.5) (1.0,  0.5)
    (-1.0,  0.0) (-0.5,  0.0) (0.0,  0.0) (0.5,  0.0) (1.0,  0.0)
    (-1.0, -0.5) (-0.5, -0.5) (0.0, -0.5) (0.5, -0.5) (1.0, -0.5)
    (-1.0, -1.0) (-0.5, -1.0) (0.0, -1.0) (0.5, -1.0) (1.0, -1.0)

    NOTE:
        The X-axis focal and Y-axis focal are assumed to be the same according
        to the pinhole camera model.

    Args:
        batch_size: Batch size of images, each of which has the same ray
            directions. Denoted as `N`.
        image_size: Size of the image. One element indicates square image, while
            two elements stand for height and width respectively. Denoted as `H`
            and `W`.
        fov: Field of view (along X axis) of the camera, in unit of degree.
        boundary: The maximum value of the X coordinate. Defaults to `1.0`.
        focal (optional): Focal Length of camera. If given, it will cover the
            focal calculated by `fov`. Note that the focal is a normalized one
            which is divided by size of the image.
        cam_look_at_dir: Direction of camera looks at. Defaults to `-1`, which
            means camera looks at `-z` direction.
        pixel_center: Whether rays originate from the pixel center or not. For
            example, assume a pixel is at (H, W). If `pixel_center` is set
            `True`, then the ray originate from (H+0.5, W+0.5), otherwise it
            originate from (H, W).
        y_descending: Whether the Y axis is in descending order from top to
            bottom. If set `True`, the coordinates are the same as the above
            example. If set `False`, the coordinate system is consistent with
            2D image plane coordinate system, where Y axis is in ascending
            order. Defaults to `True`.
    Returns:
        A tensor, with shape [N, H, W, 3], representing the per-pixel ray
            direction. Each direction is normalized to a unit vector.
    """
    # Check inputs.
    assert isinstance(batch_size, int) and batch_size > 0
    N = batch_size
    assert isinstance(image_size, (int, list, tuple))
    if isinstance(image_size, int):
        H = image_size
        W = image_size
    else:
        H, W = image_size
    assert isinstance(H, int) and H > 0
    assert isinstance(W, int) and W > 0
    assert 0 < fov < 180
    assert boundary > 0

    # Get running device.
    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'

    # Get (x, y) grid by boundary.
    max_x = boundary
    max_y = boundary / W * H
    if pixel_center:
        y, x = torch.meshgrid(
            torch.linspace(max_y - 0.5 / H, -max_y + 0.5 / H, H,
                           device=device),
            torch.linspace(-max_x + 0.5 / W, max_x - 0.5 / W, W,
                           device=device))
    else:
        y, x = torch.meshgrid(torch.linspace(max_y, -max_y, H, device=device),
                          torch.linspace(-max_x, max_x, W, device=device))
    # Get z coordinate of the image plane by focal (i.e., FOV).
    if not y_descending:
        y = -y
    if focal is None:
        focal = boundary / np.tan((2 * np.pi * fov / 360) / 2)
    z = np.sign(cam_look_at_dir) * focal * torch.ones_like(x)  # [H, W]
    # Normalize directions to unit vectors.
    rays = F.normalize(torch.stack([x, y, z], dim=-1), dim=-1)  # [H, W, 3]

    return rays.unsqueeze(0).repeat(N, 1, 1, 1)  # [N, H, W, 3]


def sample_points_per_ray(batch_size,
                          image_size,
                          num_points,
                          strategy='uniform',
                          dis_min=None,
                          dis_max=None,
                          dis_mean=None,
                          dis_stddev=None,
                          per_ray_ref=None,
                          per_point_ref=None):
    """Samples per-point radial distance on each ray.

    This function is independent of ray directions, hence, each point is
    represent by a number, indicating its radial distance to the origin (i.e.,
    the camera).

    The following sampling strategies are supported:

    - `uniform`:
        For each ray, the points uniformly locate in range `[dis_min, dis_max]`.

    - `normal`:
        For each ray, the points are sampled subject to
        `Gaussian(dis_mean, dis_stddev^2)`.

    - `ray_dependent`:
        Each ray follows a separate strategy, controlled by `per_ray_ref`.

    - `point_dependent`:
        Each point follows a separate strategy, controlled by `per_point_ref`.

    Args:
        batch_size: Batch size of images, for which points are sampled
            independently. Denoted as `N`.
        image_size: Size of the image. One element indicates square image, while
            two elements stand for height and width respectively. Denoted as `H`
            and `W`.
        num_points: Number of points sampled on each ray. Denoted as `K`.
        strategy: Strategy for point sampling. Defaults to `uniform`.
        dis_min: Minimum radial distance (with camera as the origin) for each
            point. Defaults to `None`.
        dis_max: Maximum radial distance (with camera as the origin) for each
            point. Defaults to `None`.
        dis_mean: Mean radial distance (with camera as the origin) for each
            point. Defaults to `None`.
        dis_stddev: Standard deviation of the radial distance (with camera as
            the origin) for each point. Defaults to `None`.
        per_ray_ref: Reference for each ray, which will guide the sampling
            process. Shape [N, H, W, c] is expected, where `c` is the dimension
            of a single reference. Defaults to `None`.
        per_point_ref: Reference for each point, which will guide the sampling
            process. Shape [N, H, W, K, c] is expected, where `c` is the
            dimension of a single reference. Defaults to `None`.

    Returns:
        A tensor, with shape [N, H, W, K], representing the per-point radial
            distance on each ray. All numbers should be positive, and the
            distances on each ray should follow a non-descending order.

    Raises:
        ValueError: If the sampling strategy is not supported.
        NotImplementedError: If the sampling strategy is not implemented.
    """
    # Check inputs.
    assert isinstance(batch_size, int) and batch_size > 0
    N = batch_size
    assert isinstance(image_size, (int, list, tuple))
    if isinstance(image_size, int):
        H = image_size
        W = image_size
    else:
        H, W = image_size
    assert isinstance(H, int) and H > 0
    assert isinstance(W, int) and W > 0
    assert isinstance(num_points, int) and num_points > 0
    K = num_points
    strategy = strategy.lower()
    if strategy not in _POINT_SAMPLING_STRATEGIES:
        raise ValueError(f'Invalid point sampling strategy: `{strategy}`!\n'
                         f'Strategies allowed: {_POINT_SAMPLING_STRATEGIES}.')

    # Get running device.
    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'

    # Sample points according to strategy.
    if strategy == 'uniform':
        assert dis_max >= dis_min > 0
        radii = torch.linspace(dis_min, dis_max, K, device=device)  # [K]
        return radii.reshape(1, 1, 1, K).repeat(N, H, W, 1)  # [N, H, W, K]

    if strategy == 'normal':
        # TODO: Should we support the normal sampling strategy?
        assert dis_mean > 0 and dis_stddev >= 0

    if strategy == 'ray_dependent':
        # TODO: Strategy dependent on depth?
        assert per_ray_ref.ndim == 4
        assert per_ray_ref.shape[:3] == (N, H, W)

    if strategy == 'point_dependent':
        # TODO: This is hierarchical sampling?
        assert per_point_ref.ndim == 5
        assert per_point_ref.shape[:4] == (N, H, W, K)

    raise NotImplementedError(f'Not implemented point sampling strategy: '
                              f'`{strategy}`!')


def perturb_points_per_ray(radii, strategy='middle_uniform'):
    # Stratified sampling approach described in original NeRF paper.
    """Perturbs point radii within their local range on each ray.

    `N`, `H`, `W`, `K` denote batch size, image height, image width, number of
    points per ray, respectively.

    The following perturbing strategies are supported:

    - `no`:
        Disable point perturbation.

    - `middle_uniform`:
        For each point, it is perturbed between two midpoints. One locates
        within the point itself and its previous one on the same ray, while the
        other locates within the point itself and its next one on the same ray.

    - `uniform`:
        For each point, it is perturbed between itself and its next one.
        For example, there are `n+1` points on the ray: [x_0, x_1, ..., x_n].
        Then the perturbed points are [x_0', x_1', ..., x_n'] with distribution
        xi' ~ U(x_i, x_i+1), where x_n+1 = x_n + (x_n - x_n-1).

    - `self_uniform`:
        For each point, it is perturbed around itself.For example, there are
        `n+1` points on the ray: [x_0, x_1, ..., x_n]. Then the perturbed points
        are [x_0', x_1', ..., x_n'] with distribution
        xi' ~ U(x_i - 0.5, x_i+1 - 0.5).

    Args:
        radii: A collection of point radii, with shape [N, H, W, K].
        strategy: Strategy to perturb each point. Defaults to `middle_uniform`.

    Returns:
        A tensor, with shape [N, H, W, K], representing the per-point radial
            distance on each ray. All numbers should be positive, and the
            distances on each ray should follow a non-descending order.

    Raises:
        ValueError: If the input point radii are with invalid shape, or the
            perturbing strategy is not supported.
        NotImplementedError: If the perturbing strategy is not implemented.
    """
    # Check inputs.
    if radii.ndim != 4:
        raise ValueError(f'The input point radii should be with shape '
                         f'[batch_size, height, width, num_points], '
                         f'but `{radii.shape}` is received!')
    strategy = strategy.lower()
    if strategy not in _POINT_PERTURBING_STRATEGIES:
        raise ValueError(f'Invalid point perturbing strategy: `{strategy}`!\n'
                         f'Strategies allowed: {_POINT_PERTURBING_STRATEGIES}.')

    if strategy == 'no':
        return radii

    if strategy == 'middle_uniform':
        # Get midpoints.
        midpoint = (radii[..., 1:] + radii[..., :-1]) / 2  # [N, H, W, K-1]
        # Get intervals.
        left = torch.cat([radii[..., :1], midpoint], dim=-1)  # [N, H, W, K]
        right = torch.cat([midpoint, radii[..., -1:]], dim=-1)  # [N, H, W, K]
        # Uniformly sample within each interval.
        t = torch.rand_like(radii)  # [N, H, W, K]
        return left + (right - left) * t  # [N, H, W, K]
    elif strategy == 'uniform':
        delta = radii[..., 1:2] - radii[..., 0:1]   # [N, H, W, 1]
        t = torch.rand_like(radii)   # [N, H, W, K]
        return radii + t * delta     # [N, H, W, K]
    elif strategy == 'self_uniform':
        delta = radii[..., 1:2] - radii[..., 0:1]   # [N, H, W, 1]
        t = torch.rand_like(radii) - 0.5   # [N, H, W, K]
        return radii + t * delta     # [N, H, W, K]

    raise NotImplementedError(f'Not implemented point perturbing strategy: '
                              f'`{strategy}`!')


def sample_tensor(size,
                  strategy='uniform',
                  entry_fix=None,
                  entry_min=None,
                  entry_max=None,
                  entry_mean=None,
                  entry_stddev=None):
    """Samples a tensor according to specified strategy.

    The following sampling strategies are supported:

    - `fix`:
        Each entry is fixed as `entry_fix`.

    - `uniform`:
        Each entry is uniformly sampled from range `[entry_min, entry_max]`.

    - `normal`:
        Each entry is sampled subject to `Gaussian(entry_mean, entry_stddev^2)`.

    - `hybrid`:
        Each entry is 50% sampled with `uniform` and 50% sampled with `normal`.

    - `truncated_normal`:
        Each entry is sampled subject to a truncated normal distribution, with
        `entry_min` and `entry_max` as the cut-off values.


    Args:
        size: Size of the sampled tensor. This field is expected to be an
            integer, a list, or a tuple.
        strategy: Strategy to sample points. Defaults to `uniform`.
        entry_min: Minimum value of each entry. Defaults to `None`.
        entry_max: Maximum value of each entry. Defaults to `None`.
        entry_mean: Mean value of each entry. Defaults to `None`.
        entry_stddev: Standard deviation of each entry. Defaults to `None`.

    Returns:
        A tensor, with expected size.

    Raises:
        ValueError: If the sampling strategy is not supported.
        NotImplementedError: If the sampling strategy is not implemented.
    """
    # Check inputs.
    if isinstance(size, int):
        size = (size,)
    elif isinstance(size, list):
        size = tuple(size)
    assert isinstance(size, tuple)
    strategy = strategy.lower()
    if strategy not in _TENSOR_SAMPLING_STRATEGIES:
        raise ValueError(f'Invalid tensor sampling strategy: `{strategy}`!\n'
                         f'Strategies allowed: {_TENSOR_SAMPLING_STRATEGIES}.')

    # Get running device.
    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'

    if strategy == 'fix':
        assert entry_fix is not None
        return torch.ones(size, device=device) * entry_fix

    if strategy == 'uniform':
        assert entry_max >= entry_min
        t = torch.rand(size, device=device)
        return entry_min + (entry_max - entry_min) * t

    if strategy == 'normal':
        assert entry_mean is not None and entry_stddev >= 0
        return torch.randn(size, device=device) * entry_stddev + entry_mean

    if strategy == 'hybrid':
        assert entry_max >= entry_min
        assert entry_mean is not None and entry_stddev >= 0
        if np.random.random() < 0.5:
            t = torch.rand(size, device=device)
            return entry_min + (entry_max - entry_min) * t
        return torch.randn(size, device=device) * entry_stddev + entry_mean

    if strategy == 'truncated_normal':
        # TODO: Truncated normal distribution differs from cut-off.
        assert entry_max >= entry_min
        assert entry_mean is not None and entry_stddev >= 0
        tensor = torch.randn(size, device=device) * entry_stddev + entry_mean
        tensor = torch.clamp(tensor, entry_min, entry_max)
        return tensor

    raise NotImplementedError(f'Not implemented tensor sampling strategy: '
                              f'`{strategy}`!')


def sample_camera_extrinsics(batch_size,
                             radius_strategy='fix',
                             radius_fix=None,
                             radius_min=None,
                             radius_max=None,
                             radius_mean=None,
                             radius_stddev=None,
                             polar_strategy='uniform',
                             polar_fix=None,
                             polar_min=None,
                             polar_max=None,
                             polar_mean=None,
                             polar_stddev=None,
                             azimuthal_strategy='uniform',
                             azimuthal_fix=None,
                             azimuthal_min=None,
                             azimuthal_max=None,
                             azimuthal_mean=None,
                             azimuthal_stddev=None,
                             use_spherical_uniform_position=False,
                             pitch_strategy='fix',
                             pitch_fix=0,
                             pitch_min=None,
                             pitch_max=None,
                             pitch_mean=None,
                             pitch_stddev=None,
                             yaw_strategy='fix',
                             yaw_fix=0,
                             yaw_min=None,
                             yaw_max=None,
                             yaw_mean=None,
                             yaw_stddev=None,
                             roll_strategy='fix',
                             roll_fix=0,
                             roll_min=None,
                             roll_max=None,
                             roll_mean=None,
                             roll_stddev=None):
    """Samples camera extrinsics.

    This function supports sampling camera extrinsics from 6 dimensions (here,
    all angles are in unit of radian):

    - Camera position:
        - radius: Distance from the camera position to the origin of the world
            coordinate system.
        - polar: The polar angle with respect to the origin of the world
            coordinate system.
        - azimuthal: The azimuthal angle with respect to the origin of the world
            coordinate system.
    - Camera orientation:
        - pitch: Pitch angle (X axis) regarding the camera coordinate system.
        - yaw: Yaw angle (Y axis) regarding the camera coordinate system.
        - roll: Roll angle (Z axis) regarding the camera coordinate system.

    and then convert the camera extrinsics to camera position and coordinate
    transformation matrix.

    More details about sampling as well as arguments can be found in function
    `sample_tensor()`.

    NOTE:
        Without camera orientation (i.e., `pitch = 0, yaw = 0, roll = 0`), this
        function assumes the camera pointing to the origin of the world
        coordinate system. Furthermore, camera orientation controls the rotation
        within the camera coordinate system, which is independent from the
        transformation across coordinate systems. As a result, the camera does
        not necessarily point to the origin of the world coordinate system
        anymore.

    Args:
        batch_size: Batch size of the sampled camera. Denoted as `N`.
        use_spherical_uniform_position: Whether to sample the camera position
            subject to a spherical uniform distribution. Defaults to False.

    Returns:
        A dictionary, containing
            - `camera_radius`: camera radius w.r.t. the world coordinate system,
                with shape [N].
            - `camera_polar`: camera polar w.r.t. the world coordinate system,
                with shape [N].
            - `camera_azimuthal`: camera azimuthal w.r.t. the world coordinate
                system, with shape [N].
            - `camera_pitch`: camera pitch w.r.t. the camera coordinate system,
                with shape [N].
            - `camera_yaw`: camera yaw w.r.t. the camera coordinate system,
                with shape [N].
            - `camera_roll`: camera roll w.r.t. the camera coordinate system,
                with shape [N].
            - `camera_pos`: camera position, i.e., the (x, y, z) coordinate
                in the world coordinate system, with shape [N, 3].
            - `cam2world_matrix`: transformation matrix to transform the camera
                coordinate system to the world coordinate system, with shape
                [N, 4, 4].
    """
    # Sample camera position.
    radius = sample_tensor(size=batch_size,
                           strategy=radius_strategy,
                           entry_fix=radius_fix,
                           entry_min=radius_min,
                           entry_max=radius_max,
                           entry_mean=radius_mean,
                           entry_stddev=radius_stddev)
    if use_spherical_uniform_position:
        # TODO: Check the local spherical uniform distribution?
        polar = sample_tensor(size=batch_size,
                              strategy='uniform',
                              entry_fix=polar_fix,
                              entry_min=polar_min,
                              entry_max=polar_max,
                              entry_mean=polar_mean,
                              entry_stddev=polar_stddev)
        azimuthal_cos_val = sample_tensor(size=batch_size,
                                          strategy='uniform',
                                          entry_min=azimuthal_min / np.pi,
                                          entry_max=azimuthal_max / np.pi)
        azimuthal = torch.arccos(1 - 2 * azimuthal_cos_val)
    else:
        polar = sample_tensor(size=batch_size,
                              strategy=polar_strategy,
                              entry_fix=polar_fix,
                              entry_min=polar_min,
                              entry_max=polar_max,
                              entry_mean=polar_mean,
                              entry_stddev=polar_stddev)
        azimuthal = sample_tensor(size=batch_size,
                                  strategy=azimuthal_strategy,
                                  entry_fix=azimuthal_fix,
                                  entry_min=azimuthal_min,
                                  entry_max=azimuthal_max,
                                  entry_mean=azimuthal_mean,
                                  entry_stddev=azimuthal_stddev)

    # Sample camera orientation.
    pitch = sample_tensor(size=batch_size,
                          strategy=pitch_strategy,
                          entry_fix=pitch_fix,
                          entry_min=pitch_min,
                          entry_max=pitch_max,
                          entry_mean=pitch_mean,
                          entry_stddev=pitch_stddev)
    yaw = sample_tensor(size=batch_size,
                        strategy=yaw_strategy,
                        entry_fix=yaw_fix,
                        entry_min=yaw_min,
                        entry_max=yaw_max,
                        entry_mean=yaw_mean,
                        entry_stddev=yaw_stddev)
    roll = sample_tensor(size=batch_size,
                         strategy=roll_strategy,
                         entry_fix=roll_fix,
                         entry_min=roll_min,
                         entry_max=roll_max,
                         entry_mean=roll_mean,
                         entry_stddev=roll_stddev)

    # Get running device.
    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'

    # Get camera position.
    N = batch_size
    camera_pos = torch.zeros((N, 3), device=device)
    camera_pos[:, 0] = radius * torch.sin(polar) * torch.cos(azimuthal)
    camera_pos[:, 1] = radius * torch.cos(polar)
    camera_pos[:, 2] = radius * torch.sin(polar) * torch.sin(azimuthal)

    # Get transformation matrix with the following steps.
    #   1. Use pitch, yaw, and roll to get the rotation matrix within the camera
    #      coordinate system.
    #   2. Get the forward axis, which points from the camper position to the
    #      origin of the world coordinate system.
    #   3. Get a "pseudo" up axis, which is [0, 1, 0].
    #   4. Get the left axis by crossing the "pseudo" up axis with the forward
    #      axis.
    #   5. Get the "actual" up axis by crossing the forward axis with the left
    #      axis.
    #   6. Get the camera-to-world rotation matrix with the aforementioned
    #      forward axis, left axis, and "actual" up axis.
    #   7. Get the camera-to-world transformation matrix.
    pitch_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(N, 1, 1)
    pitch_matrix[:, 1, 1] = torch.cos(pitch)
    pitch_matrix[:, 2, 2] = torch.cos(pitch)
    pitch_matrix[:, 1, 2] = -torch.sin(pitch)
    pitch_matrix[:, 2, 1] = torch.sin(pitch)  # [N, 4, 4]
    yaw_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(N, 1, 1)
    yaw_matrix[:, 0, 0] = torch.cos(yaw)
    yaw_matrix[:, 2, 2] = torch.cos(yaw)
    yaw_matrix[:, 2, 0] = -torch.sin(yaw)
    yaw_matrix[:, 0, 2] = torch.sin(yaw)  # [N, 4, 4]
    roll_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(N, 1, 1)
    roll_matrix[:, 0, 0] = torch.cos(roll)
    roll_matrix[:, 1, 1] = torch.cos(roll)
    roll_matrix[:, 0, 1] = -torch.sin(roll)
    roll_matrix[:, 1, 0] = torch.sin(roll)  # [N, 4, 4]

    forward_axis = F.normalize(camera_pos * -1, dim=-1)  # [N, 3]
    pseudo_up_axis = torch.as_tensor([0.0, 1.0, 0.0], device=device)  # [3]
    pseudo_up_axis = pseudo_up_axis.reshape(1, 3).repeat(N, 1)  # [N, 3]
    left_axis = torch.cross(pseudo_up_axis, forward_axis, dim=-1)  # [N, 3]
    left_axis = F.normalize(left_axis, dim=-1)  # [N, 3]
    up_axis = torch.cross(forward_axis, left_axis, dim=-1)  # [N, 3]
    up_axis = F.normalize(up_axis, dim=-1)  # [N, 3]

    rotation_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(N, 1, 1)
    rotation_matrix[:, :3, 0] = -left_axis
    rotation_matrix[:, :3, 1] = up_axis
    rotation_matrix[:, :3, 2] = -forward_axis  # [N, 4, 4]

    translation_matrix = torch.eye(4, device=device)
    translation_matrix = translation_matrix.unsqueeze(0).repeat(N, 1, 1)
    translation_matrix[:, :3, 3] = camera_pos  # [N, 4, 4]

    cam2world_matrix = (translation_matrix @ rotation_matrix @
                        roll_matrix @ yaw_matrix @ pitch_matrix)  # [N, 4, 4]

    return {
        'radius': radius,
        'polar': polar,
        'azimuthal': azimuthal,
        'pitch': pitch,
        'yaw': yaw,
        'roll': roll,
        'camera_pos':camera_pos,
        'cam2world_matrix': cam2world_matrix
    }


def get_point_coord(rays_camera, radii, cam2world_matrix):
    """Gets pre-point coordinate in the world coordinate system.

    `N`, `H`, `W`, `K` denote batch size, image height, image width, number of
    points per ray, respectively.

    Args:
        rays_camera: Per-pixel ray direction, with shape [N, H, W, 3], in the
            camera coordinate system.
        radii: Per-point radial distance on each ray, with shape [N, H, W, K].
        cam2world_matrix: Transformation matrix that transforms the camera
            coordinate system to the world coordinate system, with shape
            [N, 4, 4].

    Returns:
        A dictionary, containing
            - `rays_world`: ray directions in the world coordinate system,
                with shape [N, H, W, 3].
            - `points_camera`: per-point coordinate in the camera coordinate
                system, with shape [N, H, W, K, 3].
            - `points_world`: per-point coordinate in the world coordinate
                system, with shape [N, H, W, K, 3].

    Raises:
        ValueError: If any input has invalid shape.
    """
    # Check inputs.
    if rays_camera.ndim != 4 or rays_camera.shape[3] != 3:
        raise ValueError(f'The input rays should be with shape '
                         f'[batch_size, height, width, 3], '
                         f'but `{rays_camera.shape}` is received!')
    N, H, W, _ = rays_camera.shape
    if radii.ndim != 4 or radii.shape[:3] != (N, H, W):
        raise ValueError(f'The input radii should be with shape '
                         f'[batch_size, height, width, num_points],  where '
                         f'batch_size, height, width align with those of rays, '
                         f'but `{radii.shape}` is received!')
    K = radii.shape[3]
    if cam2world_matrix.shape != (N, 4, 4):
        raise ValueError(f'The input cam2world_matrix should be with shape '
                         f'[batch_size, 4, 4], where batch_size align with '
                         f'that of rays and radii '
                         f'but `{cam2world_matrix.shape}` is received!')

    # Get running device.
    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'

    # Transform rays.
    rays_world = (cam2world_matrix[:, :3, :3] @
                  rays_camera.reshape(N, -1, 3).permute(0, 2, 1))
    rays_world = rays_world.permute(0, 2, 1).reshape(N, H, W, 3)

    # Transform ray origins.
    ray_origins_homo = torch.zeros((N, H * W, 4), device=device)
    ray_origins_homo[..., 3] = 1
    ray_origins_world = torch.bmm(cam2world_matrix,
                                  ray_origins_homo.permute(0, 2, 1)).permute(
                                      0, 2, 1)[..., :3]
    ray_origins_world = ray_origins_world.reshape(N, H, W, 3)

    # Transform points.
    points_camera = (rays_camera.unsqueeze(3) *
                     radii.unsqueeze(4))  # [N, H, W, K, 3]
    points_camera_homo = torch.cat(
        [points_camera, torch.ones((N, H, W, K, 1), device=device)],
        dim=-1)  # [N, H, W, K, 4]
    points_world_homo = (cam2world_matrix @
                         points_camera_homo.reshape(N, -1, 4).permute(0, 2, 1))
    points_world = points_world_homo.permute(0, 2, 1)[:, :, :3]
    points_world = points_world.reshape(N, H, W, K, 3)

    return {
        'rays_world': rays_world,
        'ray_origins_world': ray_origins_world,
        'points_camera': points_camera,
        'points_world': points_world,
    }