# python3.8 """Contains the functions to sample points in 3D space. Point sampling is primarily used for Neural Radiance Field (NeRF). Paper: https://arxiv.org/pdf/2003.08934.pdf """ 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, # Ray sampling related. fov=30, image_boundary_value=1.0, focal=None, x_axis_right=True, y_axis_up=True, z_axis_out=True, x_pixel_shift=0.0, y_pixel_shift=0.0, selected_pixels=None, patch_params=None, # Point sampling (i.e., radial distance w.r.t. camera) related. num_points=16, point_strategy='uniform', 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()`, and `sample_camera_extrinsics()`. """ super().__init__() self.fov = fov self.image_boundary_value = image_boundary_value self.focal = focal self.x_axis_right = x_axis_right self.y_axis_up = y_axis_up self.z_axis_out = z_axis_out self.x_pixel_shift = x_pixel_shift self.y_pixel_shift = y_pixel_shift self.selected_pixels = selected_pixels self.patch_params = patch_params self.num_points = num_points self.point_strategy = point_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 def forward(self, batch_size, image_size, cam2world_matrix=None, patch_grid=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`. cam2world_matrix: Transformation matrix used to transform the camera coordinate system to the world coordinate system, with shape [N, 4, 4]. If given, the process of sampling camera extrinsics will be skipped. patch_grid: A flow-field grid is used to sample patches from the entire tensor, with shape [N, h, w, 2]. **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]. """ fov = kwargs.get('fov', self.fov) focal = kwargs.get('focal', self.focal) image_boundary_value = kwargs.get( 'image_boundary_value', self.image_boundary_value) x_axis_right = kwargs.get('x_axis_right', self.x_axis_right) y_axis_up = kwargs.get('y_axis_up', self.y_axis_up) z_axis_out = kwargs.get('z_axis_out', self.z_axis_out) x_pixel_shift = kwargs.get('x_pixel_shift', self.x_pixel_shift) y_pixel_shift = kwargs.get('y_pixel_shift', self.y_pixel_shift) selected_pixels = kwargs.get('selected_pixels', self.selected_pixels) patch_params = kwargs.get('patch_params', self.patch_params) num_points = kwargs.get('num_points', self.num_points) point_strategy = kwargs.get( 'point_strategy', self.point_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, x_axis_right=x_axis_right, y_axis_up=y_axis_up, z_axis_out=z_axis_out, x_pixel_shift=x_pixel_shift, y_pixel_shift=y_pixel_shift, selected_pixels=selected_pixels, patch_params=patch_params) if patch_grid is not None: rays_camera = rays_camera.permute(0, 3, 1, 2) rays_camera = F.grid_sample(rays_camera, patch_grid, mode='bilinear', align_corners=True) rays_camera = rays_camera.permute(0, 2, 3, 1) _, h, w, _ = patch_grid.shape image_size = h if h == w else (h, w) if selected_pixels is not None: h, w = selected_pixels.shape[1:3] image_size = h if h == w else (h, w) radii_raw = sample_points_per_ray(batch_size=batch_size, image_size=image_size, num_points=num_points, strategy=point_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) if cam2world_matrix is not None: camera_info = { 'radius': None, 'polar': None, 'azimuthal': None, 'pitch': None, 'yaw': None, 'roll': None, 'camera_pos': None, 'cam2world_matrix': cam2world_matrix, } 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] '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, x_axis_right=True, y_axis_up=True, z_axis_out=True, x_pixel_shift=0.0, y_pixel_shift=0.0, selected_pixels=None, patch_params=None, normalize=True): """Gets ray direction for each image pixel under camera coordinate system. Each ray direction is represented 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 `x_axis_right` controlling whether the positive direction points to the right hand side. - The Y axis is set as the vertical direction of the image plane, with `y_axis_up` controlling whether the positive direction points to the upside. - The Z axis is set as the direction perpendicular to the image plane, with `z_axis_out` controlling whether the positive direction points to the outside. If true, then under the camera coordinate system, the z coordinate of the image plane is negative. - By default, where `x_axis_right`, `y_axis_up`, and `z_axis_out` are all set as `True`, 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 (with `x_axis_right = True` and `y_axis_up = True`) 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) In the above case, if `x_pixel_shift = 0.5`, which means all rays will be sampled at middle-pixels along the X axis, the X shift value will be 2 * boundary / (W - 1) * 0.5 = 2 * 1 / 4 * 0.5 = 0.25 Then, the per-pixel (x, y) coordinates should look like: (-0.75, 1.0) (-0.25, 1.0) (0.25, 1.0) (0.75, 1.0) (1.25, 1.0) (-0.75, 0.5) (-0.25, 0.5) (0.25, 0.5) (0.75, 0.5) (1.25, 0.5) (-0.75, 0.0) (-0.25, 0.0) (0.25, 0.0) (0.75, 0.0) (1.25, 0.0) (-0.75, -0.5) (-0.25, -0.5) (0.25, -0.5) (0.75, -0.5) (1.25, -0.5) (-0.75, -1.0) (-0.25, -1.0) (0.25, -1.0) (0.75, -1.0) (1.25, -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: Focal Length of camera. If not given, it will be calculated by `fov` and `boundary` automatically. Note that focal is assumed to be normalized by image size. Defaults to `None`. x_axis_right: Whether the positive direction of X axis points to the right hand side. Defaults to `True`. y_axis_up: Whether the positive direction of Y axis points to the upside. Defaults to `True`. z_axis_out: Whether the positive direction of Z axis points to the outside. Defaults to `True`. x_pixel_shift: Pixel shift of each ray along X axis. Defaults to `0.0`. y_pixel_shift: Pixel shift of each ray along Y axis. Defaults to `0.0`. selected_pixels: Indices of a subset of pixels from which to sample rays, with shape [N, h, w]. If not given, all pixels will be attached a ray. Defaults to `None`. patch_params: Dictionary containing - `scales`: scales of the sampling patch, with shape [N, 2]; - `offsets`: offsets of the sampling patch, with shape [N, 2]. normalize: Whether to normalize the sampled coordinates in pixel sapce. Recall the formula of camera ray in NeRF paper : `r = o + t*d`. If `normalize` is set `True`, `t` represents the radial distance; otherwise `t` represents the depth. Returns: A tensor, with shape [N, H, W, 3] (or [N, h, w, 3] if `selected_pixels` is given), 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 if H == W: max_y = boundary else: max_y = boundary / W * H y, x = torch.meshgrid(torch.linspace(max_y, -max_y, H, device=device), torch.linspace(-max_x, max_x, W, device=device), indexing='ij') if patch_params is not None: # This patching sampling strategy is particularly used in EpiGRAF, # which includes the following steps: # 1). Shift [-1, 1] range into [0, 2]; # 2). Multiply by the patch size; # 3). Shift back to [-1, 1]; # 4). Apply the offset (converted from [0, 1] to [0, 2]). patch_scales = patch_params['scales'] patch_offsets = patch_params['offsets'] x = x.flatten().unsqueeze(0).repeat(N, 1) # [N, H * W] y = y.flatten().unsqueeze(0).repeat(N, 1) # [N, H * W] x = (x + 1.0) * patch_scales[:, 0].view( N, 1) - 1.0 + patch_offsets[:, 0].view(N, 1) * 2.0 # [N, H * W] y = (y + 1.0) * patch_scales[:, 1].view( N, 1) - 1.0 + patch_offsets[:, 1].view(N, 1) * 2.0 # [N, H * W] # Get z coordinate of the image plane by focal (i.e., FOV). if focal is None: focal = boundary / np.tan((fov / 180 * np.pi) / 2) z = -focal * torch.ones_like(x) # [H, W] # Adjust the positive direction of each axis. if not x_axis_right: x = -x if not y_axis_up: y = -y if not z_axis_out: z = -z # Adjust pixel shift along X and Y axes. x = x + 2 * max_x / (W - 1) * x_pixel_shift y = y + 2 * max_y / (H - 1) * y_pixel_shift if normalize: # Normalize directions to unit vectors. rays = F.normalize(torch.stack([x, y, z], dim=-1), dim=-1) # [H, W, 3] else: rays = torch.stack([x / focal, y / focal, z / focal], dim=-1) # [H, W, 3] # Repeat the sampled rays along the batch dimension. if rays.shape == (N, H * W, 3): rays = rays.reshape(N, H, W, 3) else: rays = rays.unsqueeze(0).repeat(N, 1, 1, 1) # [N, H, W, 3] # Select rays of interest if needed. if selected_pixels is not None: assert selected_pixels.ndim == 3 and selected_pixels.shape[0] == N _, h, w = selected_pixels.shape indices = selected_pixels.reshape(N, h * w, 1).repeat(1, 1, 3) rays = rays.reshape(N, H * W, 3) rays = torch.gather(rays, dim=1, index=indices) # [N, h * w, 3] rays = rays.reshape(N, h, w, 3) return rays 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': assert dis_mean > 0 and dis_stddev >= 0 if strategy == 'ray_dependent': assert per_ray_ref.ndim == 4 assert per_ray_ref.shape[:3] == (N, H, W) if strategy == 'point_dependent': 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] if 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] if 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_fix: Fixed value of the entry. Defaults to `None`. 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': 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, y_axis_up=True): """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. Currently, our framework supports two kind of world coordinate systems: (1) Y-axis pointing upward (default, y_upward=True): v: polar Y u: azimuth ^ |v / | / |/ +---------> X /\ / \ / u \ Z (1) Z-axis pointing upward (y_upward=False): v: polar Z u: azimuth ^ |v / | / |/ +---------> Y /\ / \ / u \ X 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 of 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: azimuthal = sample_tensor( size=batch_size, strategy='uniform', entry_min=azimuthal_min if azimuthal_min is not None else 0, entry_max=azimuthal_max if azimuthal_max is not None else 1) azimuthal = (azimuthal - 0.5) * 2 * azimuthal_stddev + azimuthal_mean polar_mean = polar_mean / np.pi polar_stddev = polar_stddev / np.pi polar = sample_tensor( size=batch_size, strategy='uniform', entry_min=polar_min if polar_min is not None else 0, entry_max=polar_max if polar_max is not None else 1) polar = (polar - 0.5) * 2 * polar_stddev + polar_mean polar = torch.arccos(1 - 2 * polar) 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) if y_axis_up: 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) else: camera_pos[:, 0] = radius * torch.sin(polar) * torch.cos(azimuthal) camera_pos[:, 1] = radius * torch.sin(polar) * torch.sin(azimuthal) camera_pos[:, 2] = radius * torch.cos(polar) # 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] if y_axis_up: pseudo_up_axis = torch.as_tensor([0.0, 1.0, 0.0], device=device) # [3] else: pseudo_up_axis = torch.as_tensor([0.0, 0.0, 1.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]. - `ray_origins_world`: ray origins 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)) ray_origins_world = ray_origins_world.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, 'points_camera': points_camera, 'points_world': points_world, }