File size: 9,479 Bytes
2568013 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 |
from math import prod
import torch
from einops import einsum, rearrange, reduce, repeat
from jaxtyping import Bool, Float, Int64
from torch import Tensor
def homogenize_points(
points: Float[Tensor, "*batch dim"],
) -> Float[Tensor, "*batch dim+1"]:
"""Convert batched points (xyz) to (xyz1)."""
return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1)
def homogenize_vectors(
vectors: Float[Tensor, "*batch dim"],
) -> Float[Tensor, "*batch dim+1"]:
"""Convert batched vectors (xyz) to (xyz0)."""
return torch.cat([vectors, torch.zeros_like(vectors[..., :1])], dim=-1)
def transform_rigid(
homogeneous_coordinates: Float[Tensor, "*#batch dim"],
transformation: Float[Tensor, "*#batch dim dim"],
) -> Float[Tensor, "*batch dim"]:
"""Apply a rigid-body transformation to points or vectors."""
return einsum(transformation, homogeneous_coordinates, "... i j, ... j -> ... i")
def transform_cam2world(
homogeneous_coordinates: Float[Tensor, "*#batch dim"],
extrinsics: Float[Tensor, "*#batch dim dim"],
) -> Float[Tensor, "*batch dim"]:
"""Transform points from 3D camera coordinates to 3D world coordinates."""
return transform_rigid(homogeneous_coordinates, extrinsics)
def transform_world2cam(
homogeneous_coordinates: Float[Tensor, "*#batch dim"],
extrinsics: Float[Tensor, "*#batch dim dim"],
) -> Float[Tensor, "*batch dim"]:
"""Transform points from 3D world coordinates to 3D camera coordinates."""
return transform_rigid(homogeneous_coordinates, extrinsics.inverse())
def project_camera_space(
points: Float[Tensor, "*#batch dim"],
intrinsics: Float[Tensor, "*#batch dim dim"],
epsilon: float = torch.finfo(torch.float32).eps,
infinity: float = 1e8,
) -> Float[Tensor, "*batch dim-1"]:
points = points / (points[..., -1:] + epsilon)
points = points.nan_to_num(posinf=infinity, neginf=-infinity)
points = einsum(intrinsics, points, "... i j, ... j -> ... i")
return points[..., :-1]
def project(
points: Float[Tensor, "*#batch dim"],
extrinsics: Float[Tensor, "*#batch dim+1 dim+1"],
intrinsics: Float[Tensor, "*#batch dim dim"],
epsilon: float = torch.finfo(torch.float32).eps,
) -> tuple[
Float[Tensor, "*batch dim-1"], # xy coordinates
Bool[Tensor, " *batch"], # whether points are in front of the camera
]:
points = homogenize_points(points)
points = transform_world2cam(points, extrinsics)[..., :-1]
in_front_of_camera = points[..., -1] >= 0
return project_camera_space(points, intrinsics, epsilon=epsilon), in_front_of_camera
def unproject(
coordinates: Float[Tensor, "*#batch dim"],
z: Float[Tensor, "*#batch"],
intrinsics: Float[Tensor, "*#batch dim+1 dim+1"],
) -> Float[Tensor, "*batch dim+1"]:
"""Unproject 2D camera coordinates with the given Z values."""
# Apply the inverse intrinsics to the coordinates.
coordinates = homogenize_points(coordinates)
ray_directions = einsum(
intrinsics.inverse(), coordinates, "... i j, ... j -> ... i"
)
# Apply the supplied depth values.
return ray_directions * z[..., None]
def get_world_rays(
coordinates: Float[Tensor, "*#batch dim"],
extrinsics: Float[Tensor, "*#batch dim+2 dim+2"],
intrinsics: Float[Tensor, "*#batch dim+1 dim+1"],
) -> tuple[
Float[Tensor, "*batch dim+1"], # origins
Float[Tensor, "*batch dim+1"], # directions
]:
# Get camera-space ray directions.
directions = unproject(
coordinates,
torch.ones_like(coordinates[..., 0]),
intrinsics,
)
directions = directions / directions.norm(dim=-1, keepdim=True)
# Transform ray directions to world coordinates.
directions = homogenize_vectors(directions)
directions = transform_cam2world(directions, extrinsics)[..., :-1]
# Tile the ray origins to have the same shape as the ray directions.
origins = extrinsics[..., :-1, -1].broadcast_to(directions.shape)
return origins, directions
def get_local_rays(
coordinates: Float[Tensor, "*#batch dim"],
intrinsics: Float[Tensor, "*#batch dim+1 dim+1"],
) -> Float[Tensor, "*batch dim+1"]:
# Get camera-space ray directions.
directions = unproject(
coordinates,
torch.ones_like(coordinates[..., 0]),
intrinsics,
)
directions = directions / directions.norm(dim=-1, keepdim=True)
return directions
def sample_image_grid(
shape: tuple[int, ...],
device: torch.device = torch.device("cpu"),
) -> tuple[
Float[Tensor, "*shape dim"], # float coordinates (xy indexing)
Int64[Tensor, "*shape dim"], # integer indices (ij indexing)
]:
"""Get normalized (range 0 to 1) coordinates and integer indices for an image."""
# Each entry is a pixel-wise integer coordinate. In the 2D case, each entry is a
# (row, col) coordinate.
indices = [torch.arange(length, device=device) for length in shape]
stacked_indices = torch.stack(torch.meshgrid(*indices, indexing="ij"), dim=-1)
# Each entry is a floating-point coordinate in the range (0, 1). In the 2D case,
# each entry is an (x, y) coordinate.
coordinates = [(idx + 0.5) / length for idx, length in zip(indices, shape)]
coordinates = reversed(coordinates)
coordinates = torch.stack(torch.meshgrid(*coordinates, indexing="xy"), dim=-1)
return coordinates, stacked_indices
def sample_training_rays(
image: Float[Tensor, "batch view channel ..."],
intrinsics: Float[Tensor, "batch view dim dim"],
extrinsics: Float[Tensor, "batch view dim+1 dim+1"],
num_rays: int,
) -> tuple[
Float[Tensor, "batch ray dim"], # origins
Float[Tensor, "batch ray dim"], # directions
Float[Tensor, "batch ray 3"], # sampled color
]:
device = extrinsics.device
b, v, _, *grid_shape = image.shape
# Generate all possible target rays.
xy, _ = sample_image_grid(tuple(grid_shape), device)
origins, directions = get_world_rays(
rearrange(xy, "... d -> ... () () d"),
extrinsics,
intrinsics,
)
origins = rearrange(origins, "... b v xy -> b (v ...) xy", b=b, v=v)
directions = rearrange(directions, "... b v xy -> b (v ...) xy", b=b, v=v)
pixels = rearrange(image, "b v c ... -> b (v ...) c")
# Sample random rays.
num_possible_rays = v * prod(grid_shape)
ray_indices = torch.randint(num_possible_rays, (b, num_rays), device=device)
batch_indices = repeat(torch.arange(b, device=device), "b -> b n", n=num_rays)
return (
origins[batch_indices, ray_indices],
directions[batch_indices, ray_indices],
pixels[batch_indices, ray_indices],
)
def intersect_rays(
origins_x: Float[Tensor, "*#batch 3"],
directions_x: Float[Tensor, "*#batch 3"],
origins_y: Float[Tensor, "*#batch 3"],
directions_y: Float[Tensor, "*#batch 3"],
eps: float = 1e-5,
inf: float = 1e10,
) -> Float[Tensor, "*batch 3"]:
"""Compute the least-squares intersection of rays. Uses the math from here:
https://math.stackexchange.com/a/1762491/286022
"""
# Broadcast the rays so their shapes match.
shape = torch.broadcast_shapes(
origins_x.shape,
directions_x.shape,
origins_y.shape,
directions_y.shape,
)
origins_x = origins_x.broadcast_to(shape)
directions_x = directions_x.broadcast_to(shape)
origins_y = origins_y.broadcast_to(shape)
directions_y = directions_y.broadcast_to(shape)
# Detect and remove batch elements where the directions are parallel.
parallel = einsum(directions_x, directions_y, "... xyz, ... xyz -> ...") > 1 - eps
origins_x = origins_x[~parallel]
directions_x = directions_x[~parallel]
origins_y = origins_y[~parallel]
directions_y = directions_y[~parallel]
# Stack the rays into (2, *shape).
origins = torch.stack([origins_x, origins_y], dim=0)
directions = torch.stack([directions_x, directions_y], dim=0)
dtype = origins.dtype
device = origins.device
# Compute n_i * n_i^T - eye(3) from the equation.
n = einsum(directions, directions, "r b i, r b j -> r b i j")
n = n - torch.eye(3, dtype=dtype, device=device).broadcast_to((2, 1, 3, 3))
# Compute the left-hand side of the equation.
lhs = reduce(n, "r b i j -> b i j", "sum")
# Compute the right-hand side of the equation.
rhs = einsum(n, origins, "r b i j, r b j -> r b i")
rhs = reduce(rhs, "r b i -> b i", "sum")
# Left-matrix-multiply both sides by the pseudo-inverse of lhs to find p.
result = torch.linalg.lstsq(lhs, rhs).solution
# Handle the case of parallel lines by setting depth to infinity.
result_all = torch.ones(shape, dtype=dtype, device=device) * inf
result_all[~parallel] = result
return result_all
def get_fov(intrinsics: Float[Tensor, "batch 3 3"]) -> Float[Tensor, "batch 2"]:
intrinsics_inv = intrinsics.inverse()
def process_vector(vector):
vector = torch.tensor(vector, dtype=torch.float32, device=intrinsics.device)
vector = einsum(intrinsics_inv, vector, "b i j, j -> b i")
return vector / vector.norm(dim=-1, keepdim=True)
left = process_vector([0, 0.5, 1])
right = process_vector([1, 0.5, 1])
top = process_vector([0.5, 0, 1])
bottom = process_vector([0.5, 1, 1])
fov_x = (left * right).sum(dim=-1).acos()
fov_y = (top * bottom).sum(dim=-1).acos()
return torch.stack((fov_x, fov_y), dim=-1)
|