Function nerf_render_view

eval_code/recons/models/moge/utils3d/torch/nerf.py:504–564

NeRF rendering of views. Note that it supports arbitrary batch dimensions (denoted as `...`) Args: extrinsics: (..., 4, 4) extrinsics matrice of the rendered views intrinsics (optional): (..., 3, 3) intrinsics matrice of the rendered views. width (optional): image w

(
    nerf: Tensor,
    extrinsics: Tensor, 
    intrinsics: Tensor, 
    width: int,
    height: int,
    *,
    patchify: bool = False,
    patch_size: Tuple[int, int] = (64, 64),
    **options: Dict[str, Any]
)

def nerf_render_view(
    nerf: Tensor,
    extrinsics: Tensor, 
    intrinsics: Tensor, 
    width: int,
    height: int,
    *,
    patchify: bool = False,
    patch_size: Tuple[int, int] = (64, 64),
    **options: Dict[str, Any]
) -> Tuple[Tensor, Tensor]:
    """
    NeRF rendering of views. Note that it supports arbitrary batch dimensions (denoted as `...`)

    Args:
        extrinsics: (..., 4, 4) extrinsics matrice of the rendered views
        intrinsics (optional): (..., 3, 3) intrinsics matrice of the rendered views.
        width (optional): image width of the rendered views.
        height (optional): image height of the rendered views.
        patchify (optional): If the image is too large, render it patch by patch
        **options: rendering options.
    
    Returns:
        rgb: (..., channels, height, width) rendered color values.
        depth: (..., height, width) rendered depth values.
    """
    if patchify:
        # Patchified rendering
        max_patch_width, max_patch_height = patch_size
        n_rows, n_columns = math.ceil(height / max_patch_height), math.ceil(width / max_patch_width)

        rgb_rows, depth_rows = [], []
        for i_row in range(n_rows):
            rgb_row, depth_row = [], []
            for i_column in range(n_columns):
                patch_shape = patch_height, patch_width = min(max_patch_height, height - i_row * max_patch_height), min(max_patch_width, width - i_column * max_patch_width)
                uv = image_uv(height, width, i_column * max_patch_width, i_row * max_patch_height, i_column * max_patch_width + patch_width, i_row * max_patch_height + patch_height).to(extrinsics)
                uv = uv.flatten(0, 1)                                               # (patch_height * patch_width, 2)
                ray_o_, ray_d_ = get_rays(extrinsics, intrinsics, uv)
                rgb_, depth_ = nerf_render_rays(nerf, ray_o_, ray_d_, **options, return_dict=False)
                rgb_ = rgb_.transpose(-1, -2).unflatten(-1, patch_shape)            # (..., 3, patch_height, patch_width)
                depth_ = depth_.unflatten(-1, patch_shape)                          # (..., patch_height, patch_width)
                
                rgb_row.append(rgb_)
                depth_row.append(depth_)
            rgb_rows.append(torch.cat(rgb_row, dim=-1))
            depth_rows.append(torch.cat(depth_row, dim=-1))
        rgb = torch.cat(rgb_rows, dim=-2)
        depth = torch.cat(depth_rows, dim=-2)

        return rgb, depth
    else:
        # Full rendering
        uv = image_uv(height, width).to(extrinsics)
        uv = uv.flatten(0, 1)                                                       # (height * width, 2)
        ray_o_, ray_d_ = get_rays(extrinsics, intrinsics, uv)
        rgb, depth = nerf_render_rays(nerf, ray_o_, ray_d_, **options, return_dict=False)
        rgb = rgb.transpose(-1, -2).unflatten(-1, (height, width))                  # (..., 3, height, width)