facebook/StyleNeRF
34
1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved2 3 4import numpy as np5from numpy.lib.function_base import angle6import torch7import torch.nn.functional as F8import math9 10from scipy.spatial.transform import Rotation as Rot11HUGE_NUMBER = 1e1012TINY_NUMBER = 1e-6 # float32 only has 7 decimal digits precision13 14 15def get_camera_mat(fov=49.13, invert=True):16 # fov = 2 * arctan(sensor / (2 * focal))17 # focal = (sensor / 2) * 1 / (tan(0.5 * fov))18 # in our case, sensor = 2 as pixels are in [-1, 1]19 focal = 1. / np.tan(0.5 * fov * np.pi/180.)20 focal = focal.astype(np.float32)21 mat = torch.tensor([22 [focal, 0., 0., 0.],23 [0., focal, 0., 0.],24 [0., 0., 1, 0.],25 [0., 0., 0., 1.]26 ]).reshape(1, 4, 4)27 if invert:28 mat = torch.inverse(mat)29 return mat30 31 32def get_random_pose(range_u, range_v, range_radius, batch_size=32,33 invert=False, gaussian=False, angular=False):34 loc, (u, v) = sample_on_sphere(range_u, range_v, size=(batch_size), gaussian=gaussian, angular=angular)35 radius = range_radius[0] + torch.rand(batch_size) * (range_radius[1] - range_radius[0])36 loc = loc * radius.unsqueeze(-1)37 R = look_at(loc)38 RT = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)39 RT[:, :3, :3] = R40 RT[:, :3, -1] = loc41 42 if invert:43 RT = torch.inverse(RT)44 45 def N(a, range_a):46 if range_a[0] == range_a[1]:47 return a * 048 return (a - range_a[0]) / (range_a[1] - range_a[0])49 50 val_u, val_v, val_r = N(u, range_u), N(v, range_v), N(radius, range_radius)51 return RT, (val_u, val_v, val_r)52 53 54def get_camera_pose(range_u, range_v, range_r, val_u=0.5, val_v=0.5, val_r=0.5,55 batch_size=32, invert=False, gaussian=False, angular=False):56 r0, rr = range_r[0], range_r[1] - range_r[0]57 r = r0 + val_r * rr58 if not gaussian:59 u0, ur = range_u[0], range_u[1] - range_u[0]60 v0, vr = range_v[0], range_v[1] - range_v[0] 61 u = u0 + val_u * ur62 v = v0 + val_v * vr63 else:64 mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 265 vu, vv = mean_u - range_u[0], mean_v - range_v[0]66 u = mean_u + vu * val_u67 v = mean_v + vv * val_v68 69 loc, _ = sample_on_sphere((u, u), (v, v), size=(batch_size), angular=angular)70 radius = torch.ones(batch_size) * r71 loc = loc * radius.unsqueeze(-1)72 R = look_at(loc)73 RT = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)74 RT[:, :3, :3] = R75 RT[:, :3, -1] = loc76 77 if invert:78 RT = torch.inverse(RT)79 return RT80 81 82def get_camera_pose_v2(range_u, range_v, range_r, mode, invert=False, gaussian=False, angular=False):83 r0, rr = range_r[0], range_r[1] - range_r[0]84 val_u, val_v = mode[:,0], mode[:,1]85 val_r = torch.ones_like(val_u) * 0.586 if not gaussian:87 u0, ur = range_u[0], range_u[1] - range_u[0]88 v0, vr = range_v[0], range_v[1] - range_v[0]89 u = u0 + val_u * ur90 v = v0 + val_v * vr91 else:92 mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 293 vu, vv = mean_u - range_u[0], mean_v - range_v[0]94 u = mean_u + vu * val_u95 v = mean_v + vv * val_v96 97 loc = to_sphere(u, v, angular)98 radius = r0 + val_r * rr99 loc = loc * radius.unsqueeze(-1)100 R = look_at(loc)101 RT = torch.eye(4).to(R.device).reshape(1, 4, 4).repeat(R.size(0), 1, 1)102 RT[:, :3, :3] = R103 RT[:, :3, -1] = loc104 105 if invert:106 RT = torch.inverse(RT)107 return RT, (val_u, val_v, val_r)108 109 110def to_sphere(u, v, angular=False):111 T = torch if isinstance(u, torch.Tensor) else np112 if not angular:113 theta = 2 * math.pi * u114 phi = T.arccos(1 - 2 * v)115 else:116 theta, phi = u, v117 118 cx = T.sin(phi) * T.cos(theta)119 cy = T.sin(phi) * T.sin(theta)120 cz = T.cos(phi)121 return T.stack([cx, cy, cz], -1)122 123 124def sample_on_sphere(range_u=(0, 1), range_v=(0, 1), size=(1,),125 to_pytorch=True, gaussian=False, angular=False):126 if not gaussian:127 u = np.random.uniform(*range_u, size=size)128 v = np.random.uniform(*range_v, size=size)129 else:130 mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 2131 var_u, var_v = mean_u - range_u[0], mean_v - range_v[0]132 u = np.random.normal(size=size) * var_u + mean_u133 v = np.random.normal(size=size) * var_v + mean_v134 135 sample = to_sphere(u, v, angular)136 if to_pytorch:137 sample = torch.tensor(sample).float()138 u, v = torch.tensor(u).float(), torch.tensor(v).float()139 140 return sample, (u, v)141 142 143def look_at(eye, at=np.array([0, 0, 0]), up=np.array([0, 0, 1]), eps=1e-5,144 to_pytorch=True):145 if not isinstance(eye, torch.Tensor):146 # this is the original code from GRAF147 at = at.astype(float).reshape(1, 3)148 up = up.astype(float).reshape(1, 3)149 eye = eye.reshape(-1, 3)150 up = up.repeat(eye.shape[0] // up.shape[0], axis=0)151 eps = np.array([eps]).reshape(1, 1).repeat(up.shape[0], axis=0)152 z_axis = eye - at153 z_axis /= np.max(np.stack([np.linalg.norm(z_axis,154 axis=1, keepdims=True), eps]))155 x_axis = np.cross(up, z_axis)156 x_axis /= np.max(np.stack([np.linalg.norm(x_axis,157 axis=1, keepdims=True), eps]))158 y_axis = np.cross(z_axis, x_axis)159 y_axis /= np.max(np.stack([np.linalg.norm(y_axis,160 axis=1, keepdims=True), eps]))161 r_mat = np.concatenate(162 (x_axis.reshape(-1, 3, 1), y_axis.reshape(-1, 3, 1), z_axis.reshape(163 -1, 3, 1)), axis=2)164 if to_pytorch:165 r_mat = torch.tensor(r_mat).float()166 else:167 168 def normalize(x, axis=-1, order=2):169 l2 = x.norm(p=order, dim=axis, keepdim=True).clamp(min=1e-8)170 return x / l2171 172 at, up = torch.from_numpy(at).float().to(eye.device), torch.from_numpy(up).float().to(eye.device)173 z_axis = normalize(eye - at[None, :])174 x_axis = normalize(torch.cross(up[None,:].expand_as(z_axis), z_axis, dim=-1))175 y_axis = normalize(torch.cross(z_axis, x_axis, dim=-1))176 r_mat = torch.stack([x_axis, y_axis, z_axis], dim=-1)177 178 return r_mat179 180 181def get_rotation_matrix(axis='z', value=0., batch_size=32):182 r = Rot.from_euler(axis, value * 2 * np.pi).as_dcm()183 r = torch.from_numpy(r).reshape(1, 3, 3).repeat(batch_size, 1, 1)184 return r185 186 187def get_corner_rays(corner_pixels, camera_matrices, res):188 assert (res + 1) * (res + 1) == corner_pixels.size(1)189 batch_size = camera_matrices[0].size(0)190 rays, origins, _ = get_camera_rays(camera_matrices, corner_pixels)191 corner_rays = torch.cat([rays, torch.cross(origins, rays, dim=-1)], -1)192 corner_rays = corner_rays.reshape(batch_size, res+1, res+1, 6).permute(0,3,1,2)193 corner_rays = torch.cat([corner_rays[..., :-1, :-1], corner_rays[..., 1:, :-1], corner_rays[..., 1:, 1:], corner_rays[..., :-1, 1:]], 1)194 return corner_rays195 196 197def arange_pixels(198 resolution=(128, 128), 199 batch_size=1, 200 subsample_to=None, 201 invert_y_axis=False, 202 margin=0,203 corner_aligned=True,204 jitter=None205 ):206 ''' Arranges pixels for given resolution in range image_range.207 208 The function returns the unscaled pixel locations as integers and the209 scaled float values.210 211 Args:212 resolution (tuple): image resolution213 batch_size (int): batch size214 subsample_to (int): if integer and > 0, the points are randomly215 subsampled to this value216 '''217 h, w = resolution218 n_points = resolution[0] * resolution[1]219 uh = 1 if corner_aligned else 1 - (1 / h)220 uw = 1 if corner_aligned else 1 - (1 / w)221 if margin > 0:222 uh = uh + (2 / h) * margin223 uw = uw + (2 / w) * margin 224 w, h = w + margin * 2, h + margin * 2225 226 x, y = torch.linspace(-uw, uw, w), torch.linspace(-uh, uh, h)227 if jitter is not None:228 dx = (torch.ones_like(x).uniform_() - 0.5) * 2 / w * jitter229 dy = (torch.ones_like(y).uniform_() - 0.5) * 2 / h * jitter230 x, y = x + dx, y + dy231 x, y = torch.meshgrid(x, y)232 pixel_scaled = torch.stack([x, y], -1).permute(1,0,2).reshape(1, -1, 2).repeat(batch_size, 1, 1)233 234 # Subsample points if subsample_to is not None and > 0235 if (subsample_to is not None and subsample_to > 0 and subsample_to < n_points):236 idx = np.random.choice(pixel_scaled.shape[1], size=(subsample_to,),237 replace=False)238 pixel_scaled = pixel_scaled[:, idx]239 240 if invert_y_axis:241 pixel_scaled[..., -1] *= -1.242 243 return pixel_scaled244 245 246def to_pytorch(tensor, return_type=False):247 ''' Converts input tensor to pytorch.248 249 Args:250 tensor (tensor): Numpy or Pytorch tensor251 return_type (bool): whether to return input type252 '''253 is_numpy = False254 if type(tensor) == np.ndarray:255 tensor = torch.from_numpy(tensor)256 is_numpy = True257 tensor = tensor.clone()258 if return_type:259 return tensor, is_numpy260 return tensor261 262 263def transform_to_world(pixels, depth, camera_mat, world_mat, scale_mat=None,264 invert=True, use_absolute_depth=True):265 ''' Transforms pixel positions p with given depth value d to world coordinates.266 267 Args:268 pixels (tensor): pixel tensor of size B x N x 2269 depth (tensor): depth tensor of size B x N x 1270 camera_mat (tensor): camera matrix271 world_mat (tensor): world matrix272 scale_mat (tensor): scale matrix273 invert (bool): whether to invert matrices (default: true)274 '''275 assert(pixels.shape[-1] == 2)276 if scale_mat is None:277 scale_mat = torch.eye(4).unsqueeze(0).repeat(278 camera_mat.shape[0], 1, 1).to(camera_mat.device)279 280 # Convert to pytorch281 pixels, is_numpy = to_pytorch(pixels, True)282 depth = to_pytorch(depth)283 camera_mat = to_pytorch(camera_mat)284 world_mat = to_pytorch(world_mat)285 scale_mat = to_pytorch(scale_mat)286 287 # Invert camera matrices288 if invert:289 camera_mat = torch.inverse(camera_mat)290 world_mat = torch.inverse(world_mat)291 scale_mat = torch.inverse(scale_mat)292 293 # Transform pixels to homogen coordinates294 pixels = pixels.permute(0, 2, 1)295 pixels = torch.cat([pixels, torch.ones_like(pixels)], dim=1)296 297 # Project pixels into camera space298 if use_absolute_depth:299 pixels[:, :2] = pixels[:, :2] * depth.permute(0, 2, 1).abs()300 pixels[:, 2:3] = pixels[:, 2:3] * depth.permute(0, 2, 1)301 else:302 pixels[:, :3] = pixels[:, :3] * depth.permute(0, 2, 1)303 304 # Transform pixels to world space305 p_world = scale_mat @ world_mat @ camera_mat @ pixels306 307 # Transform p_world back to 3D coordinates308 p_world = p_world[:, :3].permute(0, 2, 1)309 310 if is_numpy:311 p_world = p_world.numpy()312 return p_world313 314 315def transform_to_camera_space(p_world, world_mat, camera_mat=None, scale_mat=None):316 ''' Transforms world points to camera space.317 Args:318 p_world (tensor): world points tensor of size B x N x 3319 camera_mat (tensor): camera matrix320 world_mat (tensor): world matrix321 scale_mat (tensor): scale matrix322 '''323 batch_size, n_p, _ = p_world.shape324 device = p_world.device325 326 # Transform world points to homogen coordinates327 p_world = torch.cat([p_world, torch.ones(328 batch_size, n_p, 1).to(device)], dim=-1).permute(0, 2, 1)329 330 # Apply matrices to transform p_world to camera space331 if scale_mat is None:332 if camera_mat is None:333 p_cam = world_mat @ p_world334 else:335 p_cam = camera_mat @ world_mat @ p_world336 else:337 p_cam = camera_mat @ world_mat @ scale_mat @ p_world338 339 # Transform points back to 3D coordinates340 p_cam = p_cam[:, :3].permute(0, 2, 1)341 return p_cam342 343 344def origin_to_world(n_points, camera_mat, world_mat, scale_mat=None,345 invert=False):346 ''' Transforms origin (camera location) to world coordinates.347 348 Args:349 n_points (int): how often the transformed origin is repeated in the350 form (batch_size, n_points, 3)351 camera_mat (tensor): camera matrix352 world_mat (tensor): world matrix353 scale_mat (tensor): scale matrix354 invert (bool): whether to invert the matrices (default: true)355 '''356 batch_size = camera_mat.shape[0]357 device = camera_mat.device358 # Create origin in homogen coordinates359 p = torch.zeros(batch_size, 4, n_points).to(device)360 p[:, -1] = 1.361 362 if scale_mat is None:363 scale_mat = torch.eye(4).unsqueeze(364 0).repeat(batch_size, 1, 1).to(device)365 366 # Invert matrices367 if invert:368 camera_mat = torch.inverse(camera_mat)369 world_mat = torch.inverse(world_mat)370 scale_mat = torch.inverse(scale_mat)371 372 # Apply transformation373 p_world = scale_mat @ world_mat @ camera_mat @ p374 375 # Transform points back to 3D coordinates376 p_world = p_world[:, :3].permute(0, 2, 1)377 return p_world378 379 380def image_points_to_world(image_points, camera_mat, world_mat, scale_mat=None,381 invert=False, negative_depth=True):382 ''' Transforms points on image plane to world coordinates.383 384 In contrast to transform_to_world, no depth value is needed as points on385 the image plane have a fixed depth of 1.386 387 Args:388 image_points (tensor): image points tensor of size B x N x 2389 camera_mat (tensor): camera matrix390 world_mat (tensor): world matrix391 scale_mat (tensor): scale matrix392 invert (bool): whether to invert matrices393 '''394 batch_size, n_pts, dim = image_points.shape395 assert(dim == 2)396 device = image_points.device397 d_image = torch.ones(batch_size, n_pts, 1).to(device)398 if negative_depth:399 d_image *= -1.400 return transform_to_world(image_points, d_image, camera_mat, world_mat,401 scale_mat, invert=invert)402 403 404def image_points_to_camera(image_points, camera_mat, 405 invert=False, negative_depth=True, use_absolute_depth=True):406 batch_size, n_pts, dim = image_points.shape407 assert(dim == 2)408 device = image_points.device409 d_image = torch.ones(batch_size, n_pts, 1).to(device)410 if negative_depth:411 d_image *= -1.412 413 # Convert to pytorch414 pixels, is_numpy = to_pytorch(image_points, True)415 depth = to_pytorch(d_image)416 camera_mat = to_pytorch(camera_mat)417 418 # Invert camera matrices419 if invert:420 camera_mat = torch.inverse(camera_mat)421 422 # Transform pixels to homogen coordinates423 pixels = pixels.permute(0, 2, 1)424 pixels = torch.cat([pixels, torch.ones_like(pixels)], dim=1)425 426 # Project pixels into camera space427 if use_absolute_depth:428 pixels[:, :2] = pixels[:, :2] * depth.permute(0, 2, 1).abs()429 pixels[:, 2:3] = pixels[:, 2:3] * depth.permute(0, 2, 1)430 else:431 pixels[:, :3] = pixels[:, :3] * depth.permute(0, 2, 1)432 433 # Transform pixels to world space434 p_camera = camera_mat @ pixels435 436 # Transform p_world back to 3D coordinates437 p_camera = p_camera[:, :3].permute(0, 2, 1)438 439 if is_numpy:440 p_camera = p_camera.numpy()441 return p_camera442 443 444def camera_points_to_image(camera_points, camera_mat, 445 invert=False, negative_depth=True, use_absolute_depth=True):446 batch_size, n_pts, dim = camera_points.shape447 assert(dim == 3)448 device = camera_points.device449 450 # Convert to pytorch451 p_camera, is_numpy = to_pytorch(camera_points, True)452 camera_mat = to_pytorch(camera_mat)453 454 # Invert camera matrices455 if invert:456 camera_mat = torch.inverse(camera_mat)457 458 # Transform world camera space to pixels459 p_camera = p_camera.permute(0, 2, 1) # B x 3 x N460 pixels = camera_mat[:, :3, :3] @ p_camera461 462 assert use_absolute_depth and negative_depth463 pixels, p_depths = pixels[:, :2], pixels[:, 2:3]464 p_depths = -p_depths # negative depth465 pixels = pixels / p_depths466 467 pixels = pixels.permute(0, 2, 1)468 if is_numpy:469 pixels = pixels.numpy()470 return pixels471 472 473def angular_interpolation(res, camera_mat):474 batch_size = camera_mat.shape[0]475 device = camera_mat.device476 input_rays = image_points_to_camera(arange_pixels((res, res), batch_size, 477 invert_y_axis=True).to(device), camera_mat)478 output_rays = image_points_to_camera(arange_pixels((res * 2, res * 2), batch_size,479 invert_y_axis=True).to(device), camera_mat)480 input_rays = input_rays / input_rays.norm(dim=-1, keepdim=True)481 output_rays = output_rays / output_rays.norm(dim=-1, keepdim=True)482 483 def dir2sph(v):484 u = (v[..., :2] ** 2).sum(-1).sqrt()485 theta = torch.atan2(u, v[..., 2]) / math.pi486 phi = torch.atan2(v[..., 1], v[..., 0]) / math.pi487 return torch.stack([theta, phi], 1)488 489 input_rays = dir2sph(input_rays).reshape(batch_size, 2, res, res)490 output_rays = dir2sph(output_rays).reshape(batch_size, 2, res * 2, res * 2)491 return input_rays492 493 494def interpolate_sphere(z1, z2, t):495 p = (z1 * z2).sum(dim=-1, keepdim=True)496 p = p / z1.pow(2).sum(dim=-1, keepdim=True).sqrt()497 p = p / z2.pow(2).sum(dim=-1, keepdim=True).sqrt()498 omega = torch.acos(p)499 s1 = torch.sin((1-t)*omega)/torch.sin(omega)500 s2 = torch.sin(t*omega)/torch.sin(omega)501 z = s1 * z1 + s2 * z2502 return z503 504 505def get_camera_rays(camera_matrices, pixels=None, res=None, margin=0):506 device = camera_matrices[0].device507 batch_size = camera_matrices[0].shape[0]508 if pixels is None:509 assert res is not None510 pixels = arange_pixels((res, res), batch_size, invert_y_axis=True, margin=margin).to(device)511 n_points = pixels.size(1)512 pixels_world = image_points_to_world(513 pixels, camera_mat=camera_matrices[0],514 world_mat=camera_matrices[1])515 camera_world = origin_to_world(516 n_points, camera_mat=camera_matrices[0],517 world_mat=camera_matrices[1])518 ray_vector = pixels_world - camera_world519 ray_vector = ray_vector / ray_vector.norm(dim=-1, keepdim=True)520 return ray_vector, camera_world, pixels_world521 522 523def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:524 """525 Converts 6D rotation representation by Zhou et al. [1] to rotation matrix526 using Gram--Schmidt orthogonalization per Section B of [1].527 Args:528 d6: 6D rotation representation, of size (*, 6)529 530 Returns:531 batch of rotation matrices of size (*, 3, 3)532 533 [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.534 On the Continuity of Rotation Representations in Neural Networks.535 IEEE Conference on Computer Vision and Pattern Recognition, 2019.536 Retrieved from http://arxiv.org/abs/1812.07035537 """538 539 a1, a2 = d6[..., :3], d6[..., 3:]540 b1 = F.normalize(a1, dim=-1)541 b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1542 b2 = F.normalize(b2, dim=-1)543 b3 = torch.cross(b1, b2, dim=-1)544 return torch.stack((b1, b2, b3), dim=-2)545 546 547def camera_9d_to_16d(d9):548 d6, translation = d9[..., :6], d9[..., 6:]549 rotation = rotation_6d_to_matrix(d6)550 RT = torch.eye(4).to(device=d9.device, dtype=d9.dtype).reshape(551 1, 4, 4).repeat(d6.size(0), 1, 1)552 RT[:, :3, :3] = rotation553 RT[:, :3, -1] = translation554 return RT.reshape(-1, 16)555 556def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:557 """558 Converts rotation matrices to 6D rotation representation by Zhou et al. [1]559 by dropping the last row. Note that 6D representation is not unique.560 Args:561 matrix: batch of rotation matrices of size (*, 3, 3)562 563 Returns:564 6D rotation representation, of size (*, 6)565 566 [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.567 On the Continuity of Rotation Representations in Neural Networks.568 IEEE Conference on Computer Vision and Pattern Recognition, 2019.569 Retrieved from http://arxiv.org/abs/1812.07035570 """571 return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)572 573 574def depth2pts_outside(ray_o, ray_d, depth):575 '''576 ray_o, ray_d: [..., 3]577 depth: [...]; inverse of distance to sphere origin578 '''579 # note: d1 becomes negative if this mid point is behind camera580 d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)581 p_mid = ray_o + d1.unsqueeze(-1) * ray_d582 p_mid_norm = torch.norm(p_mid, dim=-1)583 ray_d_cos = 1. / torch.norm(ray_d, dim=-1)584 d2 = torch.sqrt(1. - p_mid_norm * p_mid_norm) * ray_d_cos585 p_sphere = ray_o + (d1 + d2).unsqueeze(-1) * ray_d586 587 rot_axis = torch.cross(ray_o, p_sphere, dim=-1)588 rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)589 phi = torch.asin(p_mid_norm)590 theta = torch.asin(p_mid_norm * depth) # depth is inside [0, 1]591 rot_angle = (phi - theta).unsqueeze(-1) # [..., 1]592 593 # now rotate p_sphere594 # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula595 p_sphere_new = p_sphere * torch.cos(rot_angle) + \596 torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \597 rot_axis * torch.sum(rot_axis*p_sphere, dim=-1, keepdim=True) * (1.-torch.cos(rot_angle))598 p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)599 pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)600 601 # now calculate conventional depth602 depth_real = 1. / (depth + TINY_NUMBER) * torch.cos(theta) * ray_d_cos + d1603 return pts, depth_real604 605 606def intersect_sphere(ray_o, ray_d, radius=1):607 '''608 ray_o, ray_d: [..., 3]609 compute the depth of the intersection point between this ray and unit sphere610 '''611 # note: d1 becomes negative if this mid point is behind camera612 d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)613 p = ray_o + d1.unsqueeze(-1) * ray_d614 # consider the case where the ray does not intersect the sphere615 ray_d_cos = 1. / torch.norm(ray_d, dim=-1)616 d2 = radius ** 2 - torch.sum(p * p, dim=-1)617 mask = (d2 > 0)618 d2 = torch.sqrt(d2.clamp(min=1e-6)) * ray_d_cos619 d1, d2 = d1.unsqueeze(-1), d2.unsqueeze(-1)620 depth_range = [d1 - d2, d1 + d2]621 return depth_range, mask622 623 624def normalize(x, axis=-1, order=2):625 if isinstance(x, torch.Tensor):626 l2 = x.norm(p=order, dim=axis, keepdim=True)627 return x / (l2 + 1e-8), l2628 629 else:630 l2 = np.linalg.norm(x, order, axis)631 l2 = np.expand_dims(l2, axis)632 l2[l2==0] = 1633 return x / l2, l2634 635 636def sample_pdf(bins, weights, N_importance, det=False, eps=1e-5):637 """638 Sample @N_importance samples from @bins with distribution defined by @weights.639 Inputs:640 bins: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2"641 weights: (N_rays, N_samples_)642 N_importance: the number of samples to draw from the distribution643 det: deterministic or not644 eps: a small number to prevent division by zero645 Outputs:646 samples: the sampled samples647 Source: https://github.com/kwea123/nerf_pl/blob/master/models/rendering.py648 """649 N_rays, N_samples_ = weights.shape650 weights = weights + eps # prevent division by zero (don't do inplace op!)651 pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_)652 cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function653 cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1) # (N_rays, N_samples_+1)654 # padded to 0~1 inclusive655 656 if det:657 u = torch.linspace(0, 1, N_importance, device=bins.device)658 u = u.expand(N_rays, N_importance)659 else:660 u = torch.rand(N_rays, N_importance, device=bins.device)661 u = u.contiguous()662 663 inds = torch.searchsorted(cdf, u)664 below = torch.clamp_min(inds-1, 0)665 above = torch.clamp_max(inds, N_samples_)666 667 inds_sampled = torch.stack([below, above], -1).view(N_rays, 2*N_importance)668 cdf_g = torch.gather(cdf, 1, inds_sampled)669 cdf_g = cdf_g.view(N_rays, N_importance, 2)670 bins_g = torch.gather(bins, 1, inds_sampled).view(N_rays, N_importance, 2)671 672 denom = cdf_g[...,1]-cdf_g[...,0]673 denom[denom<eps] = 1 # denom equals 0 means a bin has weight 0, in which case it will not be sampled674 # anyway, therefore any value for it is fine (set to 1 here)675 676 samples = bins_g[...,0] + (u-cdf_g[...,0])/denom * (bins_g[...,1]-bins_g[...,0])677 return samples678 679 680def normalization_inverse_sqrt_dist_centered(x_in_world, view_cell_center, max_depth):681 localized = x_in_world - view_cell_center682 local = torch.sqrt(torch.linalg.norm(localized, dim=-1))683 res = localized / (math.sqrt(max_depth) * local[..., None])684 return res685 686 687######################################################################################688