BerfScene / models /rendering /point_sampler.py
3v324v23's picture
init
2f85de4
raw
history blame
45.9 kB
# 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,
}