|
|
import torch.nn as nn |
|
|
import torch |
|
|
import numpy as np |
|
|
from .embedders import get_embedder |
|
|
|
|
|
class ImplicitNet(nn.Module): |
|
|
def __init__(self, opt): |
|
|
super().__init__() |
|
|
|
|
|
dims = [opt.d_in] + list( |
|
|
opt.dims) + [opt.d_out + opt.feature_vector_size] |
|
|
self.num_layers = len(dims) |
|
|
self.skip_in = opt.skip_in |
|
|
self.embed_fn = None |
|
|
self.opt = opt |
|
|
|
|
|
if opt.multires > 0: |
|
|
embed_fn, input_ch = get_embedder(opt.multires, input_dims=opt.d_in, mode=opt.embedder_mode) |
|
|
self.embed_fn = embed_fn |
|
|
dims[0] = input_ch |
|
|
self.cond = opt.cond |
|
|
if self.cond == 'smpl': |
|
|
self.cond_layer = [0] |
|
|
self.cond_dim = 69 |
|
|
elif self.cond == 'frame': |
|
|
self.cond_layer = [0] |
|
|
self.cond_dim = opt.dim_frame_encoding |
|
|
self.dim_pose_embed = 0 |
|
|
if self.dim_pose_embed > 0: |
|
|
self.lin_p0 = nn.Linear(self.cond_dim, self.dim_pose_embed) |
|
|
self.cond_dim = self.dim_pose_embed |
|
|
for l in range(0, self.num_layers - 1): |
|
|
if l + 1 in self.skip_in: |
|
|
out_dim = dims[l + 1] - dims[0] |
|
|
else: |
|
|
out_dim = dims[l + 1] |
|
|
|
|
|
if self.cond != 'none' and l in self.cond_layer: |
|
|
lin = nn.Linear(dims[l] + self.cond_dim, out_dim) |
|
|
else: |
|
|
lin = nn.Linear(dims[l], out_dim) |
|
|
if opt.init == 'geometry': |
|
|
if l == self.num_layers - 2: |
|
|
torch.nn.init.normal_(lin.weight, |
|
|
mean=np.sqrt(np.pi) / |
|
|
np.sqrt(dims[l]), |
|
|
std=0.0001) |
|
|
torch.nn.init.constant_(lin.bias, -opt.bias) |
|
|
elif opt.multires > 0 and l == 0: |
|
|
torch.nn.init.constant_(lin.bias, 0.0) |
|
|
torch.nn.init.constant_(lin.weight[:, 3:], 0.0) |
|
|
torch.nn.init.normal_(lin.weight[:, :3], 0.0, |
|
|
np.sqrt(2) / np.sqrt(out_dim)) |
|
|
elif opt.multires > 0 and l in self.skip_in: |
|
|
torch.nn.init.constant_(lin.bias, 0.0) |
|
|
torch.nn.init.normal_(lin.weight, 0.0, |
|
|
np.sqrt(2) / np.sqrt(out_dim)) |
|
|
torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], |
|
|
0.0) |
|
|
else: |
|
|
torch.nn.init.constant_(lin.bias, 0.0) |
|
|
torch.nn.init.normal_(lin.weight, 0.0, |
|
|
np.sqrt(2) / np.sqrt(out_dim)) |
|
|
if opt.init == 'zero': |
|
|
init_val = 1e-5 |
|
|
if l == self.num_layers - 2: |
|
|
torch.nn.init.constant_(lin.bias, 0.0) |
|
|
torch.nn.init.uniform_(lin.weight, -init_val, init_val) |
|
|
if opt.weight_norm: |
|
|
lin = nn.utils.weight_norm(lin) |
|
|
setattr(self, "lin" + str(l), lin) |
|
|
self.softplus = nn.Softplus(beta=100) |
|
|
|
|
|
def forward(self, input, cond, current_epoch=None): |
|
|
if input.ndim == 2: input = input.unsqueeze(0) |
|
|
|
|
|
num_batch, num_point, num_dim = input.shape |
|
|
|
|
|
if num_batch * num_point == 0: return input |
|
|
|
|
|
input = input.reshape(num_batch * num_point, num_dim) |
|
|
|
|
|
if self.cond != 'none': |
|
|
num_batch, num_cond = cond[self.cond].shape |
|
|
|
|
|
input_cond = cond[self.cond].unsqueeze(1).expand(num_batch, num_point, num_cond) |
|
|
|
|
|
input_cond = input_cond.reshape(num_batch * num_point, num_cond) |
|
|
|
|
|
if self.dim_pose_embed: |
|
|
input_cond = self.lin_p0(input_cond) |
|
|
|
|
|
if self.embed_fn is not None: |
|
|
input = self.embed_fn(input) |
|
|
|
|
|
x = input |
|
|
|
|
|
for l in range(0, self.num_layers - 1): |
|
|
lin = getattr(self, "lin" + str(l)) |
|
|
if self.cond != 'none' and l in self.cond_layer: |
|
|
x = torch.cat([x, input_cond], dim=-1) |
|
|
if l in self.skip_in: |
|
|
x = torch.cat([x, input], 1) / np.sqrt(2) |
|
|
x = lin(x) |
|
|
if l < self.num_layers - 2: |
|
|
x = self.softplus(x) |
|
|
|
|
|
x = x.reshape(num_batch, num_point, -1) |
|
|
|
|
|
return x |
|
|
|
|
|
def gradient(self, x, cond): |
|
|
x.requires_grad_(True) |
|
|
y = self.forward(x, cond)[:, :1] |
|
|
d_output = torch.ones_like(y, requires_grad=False, device=y.device) |
|
|
gradients = torch.autograd.grad(outputs=y, |
|
|
inputs=x, |
|
|
grad_outputs=d_output, |
|
|
create_graph=True, |
|
|
retain_graph=True, |
|
|
only_inputs=True)[0] |
|
|
return gradients.unsqueeze(1) |
|
|
|
|
|
|
|
|
class RenderingNet(nn.Module): |
|
|
def __init__(self, opt): |
|
|
super().__init__() |
|
|
|
|
|
self.mode = opt.mode |
|
|
dims = [opt.d_in + opt.feature_vector_size] + list( |
|
|
opt.dims) + [opt.d_out] |
|
|
|
|
|
self.embedview_fn = None |
|
|
if opt.multires_view > 0: |
|
|
embedview_fn, input_ch = get_embedder(opt.multires_view) |
|
|
self.embedview_fn = embedview_fn |
|
|
dims[0] += (input_ch - 3) |
|
|
if self.mode == 'nerf_frame_encoding': |
|
|
dims[0] += opt.dim_frame_encoding |
|
|
if self.mode == 'pose': |
|
|
self.dim_cond_embed = 8 |
|
|
self.cond_dim = 69 |
|
|
|
|
|
self.lin_pose = torch.nn.Linear(self.cond_dim, self.dim_cond_embed) |
|
|
self.num_layers = len(dims) |
|
|
for l in range(0, self.num_layers - 1): |
|
|
out_dim = dims[l + 1] |
|
|
lin = nn.Linear(dims[l], out_dim) |
|
|
if opt.weight_norm: |
|
|
lin = nn.utils.weight_norm(lin) |
|
|
setattr(self, "lin" + str(l), lin) |
|
|
self.relu = nn.ReLU() |
|
|
self.sigmoid = nn.Sigmoid() |
|
|
|
|
|
def forward(self, points, normals, view_dirs, body_pose, feature_vectors, frame_latent_code=None): |
|
|
if self.embedview_fn is not None: |
|
|
if self.mode == 'nerf_frame_encoding': |
|
|
view_dirs = self.embedview_fn(view_dirs) |
|
|
|
|
|
if self.mode == 'nerf_frame_encoding': |
|
|
frame_latent_code = frame_latent_code.expand(view_dirs.shape[0], -1) |
|
|
rendering_input = torch.cat([view_dirs, frame_latent_code, feature_vectors], dim=-1) |
|
|
elif self.mode == 'pose': |
|
|
num_points = points.shape[0] |
|
|
body_pose = body_pose.unsqueeze(1).expand(-1, num_points, -1).reshape(num_points, -1) |
|
|
body_pose = self.lin_pose(body_pose) |
|
|
rendering_input = torch.cat([points, normals, body_pose, feature_vectors], dim=-1) |
|
|
else: |
|
|
raise NotImplementedError |
|
|
|
|
|
x = rendering_input |
|
|
for l in range(0, self.num_layers - 1): |
|
|
lin = getattr(self, "lin" + str(l)) |
|
|
x = lin(x) |
|
|
if l < self.num_layers - 2: |
|
|
x = self.relu(x) |
|
|
x = self.sigmoid(x) |
|
|
return x |
|
|
|
