# ORIGINAL LICENSE
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Modified by Jiale Xu
# The modifications are subject to the same license as the original.
import itertools
import torch
import torch.nn as nn
from .utils.renderer import ImportanceRenderer
from .utils.ray_sampler import RaySampler
class OSGDecoder(nn.Module):
"""
Triplane decoder that gives RGB and sigma values from sampled features.
Using ReLU here instead of Softplus in the original implementation.
Reference:
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L112
"""
def __init__(self, n_features: int,
hidden_dim: int = 64, num_layers: int = 4, activation: nn.Module = nn.ReLU):
super().__init__()
self.net = nn.Sequential(
nn.Linear(3 * n_features, hidden_dim),
activation(),
*itertools.chain(*[[
nn.Linear(hidden_dim, hidden_dim),
activation(),
] for _ in range(num_layers - 2)]),
nn.Linear(hidden_dim, 1 + 3),
)
# init all bias to zero
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.zeros_(m.bias)
def forward(self, sampled_features, ray_directions):
# Aggregate features by mean
# sampled_features = sampled_features.mean(1)
# Aggregate features by concatenation
_N, n_planes, _M, _C = sampled_features.shape
sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
x = sampled_features
N, M, C = x.shape
x = x.contiguous().view(N*M, C)
x = self.net(x)
x = x.view(N, M, -1)
rgb = torch.sigmoid(x[..., 1:])*(1 + 2*0.001) - 0.001 # Uses sigmoid clamping from MipNeRF
sigma = x[..., 0:1]
return {'rgb': rgb, 'sigma': sigma}
class TriplaneSynthesizer(nn.Module):
"""
Synthesizer that renders a triplane volume with planes and a camera.
Reference:
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L19
"""
DEFAULT_RENDERING_KWARGS = {
'ray_start': 'auto',
'ray_end': 'auto',
'box_warp': 2.,
'white_back': True,
'disparity_space_sampling': False,
'clamp_mode': 'softplus',
'sampler_bbox_min': -1.,
'sampler_bbox_max': 1.,
}
def __init__(self, triplane_dim: int, samples_per_ray: int):
super().__init__()
# attributes
self.triplane_dim = triplane_dim
self.rendering_kwargs = {
**self.DEFAULT_RENDERING_KWARGS,
'depth_resolution': samples_per_ray // 2,
'depth_resolution_importance': samples_per_ray // 2,
}
# renderings
self.renderer = ImportanceRenderer()
self.ray_sampler = RaySampler()
# modules
self.decoder = OSGDecoder(n_features=triplane_dim)
def forward(self, planes, cameras, render_size=128, crop_params=None):
# planes: (N, 3, D', H', W')
# cameras: (N, M, D_cam)
# render_size: int
assert planes.shape[0] == cameras.shape[0], "Batch size mismatch for planes and cameras"
N, M = cameras.shape[:2]
cam2world_matrix = cameras[..., :16].view(N, M, 4, 4)
intrinsics = cameras[..., 16:25].view(N, M, 3, 3)
# Create a batch of rays for volume rendering
ray_origins, ray_directions = self.ray_sampler(
cam2world_matrix=cam2world_matrix.reshape(-1, 4, 4),
intrinsics=intrinsics.reshape(-1, 3, 3),
render_size=render_size,
)
assert N*M == ray_origins.shape[0], "Batch size mismatch for ray_origins"
assert ray_origins.dim() == 3, "ray_origins should be 3-dimensional"
# Crop rays if crop_params is available
if crop_params is not None:
ray_origins = ray_origins.reshape(N*M, render_size, render_size, 3)
ray_directions = ray_directions.reshape(N*M, render_size, render_size, 3)
i, j, h, w = crop_params
ray_origins = ray_origins[:, i:i+h, j:j+w, :].reshape(N*M, -1, 3)
ray_directions = ray_directions[:, i:i+h, j:j+w, :].reshape(N*M, -1, 3)
# Perform volume rendering
rgb_samples, depth_samples, weights_samples = self.renderer(
planes.repeat_interleave(M, dim=0), self.decoder, ray_origins, ray_directions, self.rendering_kwargs,
)
# Reshape into 'raw' neural-rendered image
if crop_params is not None:
Himg, Wimg = crop_params[2:]
else:
Himg = Wimg = render_size
rgb_images = rgb_samples.permute(0, 2, 1).reshape(N, M, rgb_samples.shape[-1], Himg, Wimg).contiguous()
depth_images = depth_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
weight_images = weights_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
out = {
'images_rgb': rgb_images,
'images_depth': depth_images,
'images_weight': weight_images,
}
return out
def forward_grid(self, planes, grid_size: int, aabb: torch.Tensor = None):
# planes: (N, 3, D', H', W')
# grid_size: int
# aabb: (N, 2, 3)
if aabb is None:
aabb = torch.tensor([
[self.rendering_kwargs['sampler_bbox_min']] * 3,
[self.rendering_kwargs['sampler_bbox_max']] * 3,
], device=planes.device, dtype=planes.dtype).unsqueeze(0).repeat(planes.shape[0], 1, 1)
assert planes.shape[0] == aabb.shape[0], "Batch size mismatch for planes and aabb"
N = planes.shape[0]
# create grid points for triplane query
grid_points = []
for i in range(N):
grid_points.append(torch.stack(torch.meshgrid(
torch.linspace(aabb[i, 0, 0], aabb[i, 1, 0], grid_size, device=planes.device),
torch.linspace(aabb[i, 0, 1], aabb[i, 1, 1], grid_size, device=planes.device),
torch.linspace(aabb[i, 0, 2], aabb[i, 1, 2], grid_size, device=planes.device),
indexing='ij',
), dim=-1).reshape(-1, 3))
cube_grid = torch.stack(grid_points, dim=0).to(planes.device)
features = self.forward_points(planes, cube_grid)
# reshape into grid
features = {
k: v.reshape(N, grid_size, grid_size, grid_size, -1)
for k, v in features.items()
}
return features
def forward_points(self, planes, points: torch.Tensor, chunk_size: int = 2**20):
# planes: (N, 3, D', H', W')
# points: (N, P, 3)
N, P = points.shape[:2]
# query triplane in chunks
outs = []
for i in range(0, points.shape[1], chunk_size):
chunk_points = points[:, i:i+chunk_size]
# query triplane
chunk_out = self.renderer.run_model_activated(
planes=planes,
decoder=self.decoder,
sample_coordinates=chunk_points,
sample_directions=torch.zeros_like(chunk_points),
options=self.rendering_kwargs,
)
outs.append(chunk_out)
# concatenate the outputs
point_features = {
k: torch.cat([out[k] for out in outs], dim=1)
for k in outs[0].keys()
}
return point_features