NeRF rendering of rays. Note that it supports arbitrary batch dimensions (denoted as `...`) Args: nerf: nerf model, which takes (points, directions) as input and returns (color, density) as output. If nerf is a tuple, it should be (nerf_coarse, nerf_fine), where nerf_co
(
nerf: Union[Callable[[Tensor, Tensor], Tuple[Tensor, Tensor]], Tuple[Callable[[Tensor], Tuple[Tensor, Tensor]], Callable[[Tensor], Tuple[Tensor, Tensor]]]],
rays_o: Tensor, rays_d: Tensor,
*,
return_dict: bool = False,
n_coarse: int = 64, n_fine: int = 64,
near: float = 0.1, far: float = 100.0,
z_spacing: Literal['linear', 'inverse_linear'] = 'linear',
)
| 232 | |
| 233 | |
| 234 | def nerf_render_rays( |
| 235 | nerf: Union[Callable[[Tensor, Tensor], Tuple[Tensor, Tensor]], Tuple[Callable[[Tensor], Tuple[Tensor, Tensor]], Callable[[Tensor], Tuple[Tensor, Tensor]]]], |
| 236 | rays_o: Tensor, rays_d: Tensor, |
| 237 | *, |
| 238 | return_dict: bool = False, |
| 239 | n_coarse: int = 64, n_fine: int = 64, |
| 240 | near: float = 0.1, far: float = 100.0, |
| 241 | z_spacing: Literal['linear', 'inverse_linear'] = 'linear', |
| 242 | ): |
| 243 | """ |
| 244 | NeRF rendering of rays. Note that it supports arbitrary batch dimensions (denoted as `...`) |
| 245 | |
| 246 | Args: |
| 247 | nerf: nerf model, which takes (points, directions) as input and returns (color, density) as output. |
| 248 | If nerf is a tuple, it should be (nerf_coarse, nerf_fine), where nerf_coarse and nerf_fine are two nerf models for coarse and fine stages respectively. |
| 249 | |
| 250 | nerf args: |
| 251 | points: (..., n_rays, n_samples, 3) |
| 252 | directions: (..., n_rays, n_samples, 3) |
| 253 | nerf returns: |
| 254 | color: (..., n_rays, n_samples, 3) color values. |
| 255 | density: (..., n_rays, n_samples) density values. |
| 256 | |
| 257 | rays_o: (..., n_rays, 3) ray origins |
| 258 | rays_d: (..., n_rays, 3) ray directions. |
| 259 | pixel_width: (..., n_rays) pixel width. How to compute? pixel_width = 1 / (normalized focal length * width) |
| 260 | |
| 261 | Returns |
| 262 | if return_dict is False, return rendered rgb and depth for short cut. (If there are separate coarse and fine results, return fine results) |
| 263 | rgb: (..., n_rays, 3) rendered color values. |
| 264 | depth: (..., n_rays) rendered depth values. |
| 265 | else, return a dict. If `n_fine == 0` or `nerf` is a single model, the dict only contains coarse results: |
| 266 | ``` |
| 267 | {'rgb': .., 'depth': .., 'weights': .., 'z_vals': .., 'color': .., 'density': ..} |
| 268 | ``` |
| 269 | If there are two models for coarse and fine stages, the dict contains both coarse and fine results: |
| 270 | ``` |
| 271 | { |
| 272 | "coarse": {'rgb': .., 'depth': .., 'weights': .., 'z_vals': .., 'color': .., 'density': ..}, |
| 273 | "fine": {'rgb': .., 'depth': .., 'weights': .., 'z_vals': .., 'color': .., 'density': ..} |
| 274 | } |
| 275 | ``` |
| 276 | """ |
| 277 | if isinstance(nerf, tuple): |
| 278 | nerf_coarse, nerf_fine = nerf |
| 279 | else: |
| 280 | nerf_coarse = nerf_fine = nerf |
| 281 | # 1. Coarse: bin sampling |
| 282 | z_coarse = bin_sample(rays_d.shape[:-1], n_coarse, near, far, device=rays_o.device, dtype=rays_o.dtype, spacing=z_spacing) # (n_batch, n_views, n_rays, n_samples) |
| 283 | points_coarse = rays_o[..., None, :] + rays_d[..., None, :] * z_coarse[..., None] # (n_batch, n_views, n_rays, n_samples, 3) |
| 284 | ray_length = rays_d.norm(dim=-1) |
| 285 | |
| 286 | # Query color and density |
| 287 | color_coarse, density_coarse = nerf_coarse(points_coarse, rays_d[..., None, :].expand_as(points_coarse)) # (n_batch, n_views, n_rays, n_samples, 3), (n_batch, n_views, n_rays, n_samples) |
| 288 | |
| 289 | # Volume rendering |
| 290 | with torch.no_grad(): |
| 291 | rgb_coarse, depth_coarse, weights = volume_rendering(color_coarse, density_coarse, z_coarse, ray_length) # (n_batch, n_views, n_rays, 3), (n_batch, n_views, n_rays, 1), (n_batch, n_views, n_rays, n_samples) |
no test coverage detected