392 lines
12 KiB
Python
392 lines
12 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# ------------------------------------------------------------------------
|
|
# Modified from https://github.com/Totoro97/NeuS/blob/main/models/fields.py
|
|
# Copyright (c) 2021 Peng Wang. All Rights Reserved.
|
|
# ------------------------------------------------------------------------
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.nn.utils import spectral_norm
|
|
|
|
|
|
class SDFNetwork(nn.Module):
|
|
|
|
def __init__(self,
|
|
d_in,
|
|
d_out,
|
|
d_hidden,
|
|
n_layers,
|
|
skip_in=(4, ),
|
|
multires=0,
|
|
bias=0.5,
|
|
scale=1,
|
|
geometric_init=True,
|
|
weight_norm=True,
|
|
inside_outside=False):
|
|
super(SDFNetwork, self).__init__()
|
|
|
|
dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
|
|
|
|
self.embed_fn_fine = None
|
|
|
|
if multires > 0:
|
|
embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
|
|
self.embed_fn_fine = embed_fn
|
|
dims[0] = input_ch
|
|
|
|
self.num_layers = len(dims)
|
|
self.skip_in = skip_in
|
|
self.scale = scale
|
|
|
|
for layer in range(0, self.num_layers - 1):
|
|
if layer + 1 in self.skip_in:
|
|
out_dim = dims[layer + 1] - dims[0]
|
|
else:
|
|
out_dim = dims[layer + 1]
|
|
|
|
lin = nn.Linear(dims[layer], out_dim)
|
|
|
|
if geometric_init:
|
|
if layer == self.num_layers - 2:
|
|
if not inside_outside:
|
|
torch.nn.init.normal_(
|
|
lin.weight,
|
|
mean=np.sqrt(np.pi) / np.sqrt(dims[layer]),
|
|
std=0.0001)
|
|
torch.nn.init.constant_(lin.bias, -bias)
|
|
else:
|
|
torch.nn.init.normal_(
|
|
lin.weight,
|
|
mean=-np.sqrt(np.pi) / np.sqrt(dims[layer]),
|
|
std=0.0001)
|
|
torch.nn.init.constant_(lin.bias, bias)
|
|
elif multires > 0 and layer == 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 multires > 0 and layer 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 weight_norm:
|
|
lin = nn.utils.weight_norm(lin)
|
|
|
|
setattr(self, 'lin' + str(layer), lin)
|
|
|
|
self.activation = nn.Softplus(beta=100)
|
|
|
|
def forward(self, inputs):
|
|
inputs = inputs * self.scale
|
|
|
|
if self.embed_fn_fine is not None:
|
|
inputs = self.embed_fn_fine(inputs)
|
|
|
|
x = inputs
|
|
for layer in range(0, self.num_layers - 1):
|
|
lin = getattr(self, 'lin' + str(layer))
|
|
|
|
if layer in self.skip_in:
|
|
x = torch.cat([x, inputs], 1) / np.sqrt(2)
|
|
|
|
x = lin(x)
|
|
|
|
if layer < self.num_layers - 2:
|
|
x = self.activation(x)
|
|
return torch.cat([x[:, :1] / self.scale, x[:, 1:]], dim=-1)
|
|
|
|
def sdf(self, x):
|
|
return self.forward(x)[:, :1]
|
|
|
|
def sdf_hidden_appearance(self, x):
|
|
return self.forward(x)
|
|
|
|
def gradient(self, x):
|
|
x.requires_grad_(True)
|
|
with torch.enable_grad():
|
|
y = self.sdf(x)
|
|
|
|
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 RenderingNetwork(nn.Module):
|
|
|
|
def __init__(self,
|
|
d_feature,
|
|
mode,
|
|
d_in,
|
|
d_out,
|
|
d_hidden,
|
|
n_layers,
|
|
weight_norm=True,
|
|
multires_view=0,
|
|
squeeze_out=True):
|
|
super().__init__()
|
|
|
|
self.mode = mode
|
|
self.squeeze_out = squeeze_out
|
|
dims = [d_in + d_feature] + [d_hidden
|
|
for _ in range(n_layers)] + [d_out]
|
|
|
|
self.embedview_fn = None
|
|
if multires_view > 0:
|
|
embedview_fn, input_ch = get_embedder(multires_view)
|
|
self.embedview_fn = embedview_fn
|
|
dims[0] += (input_ch - 3)
|
|
|
|
self.num_layers = len(dims)
|
|
|
|
for layer in range(0, self.num_layers - 1):
|
|
out_dim = dims[layer + 1]
|
|
lin = nn.Linear(dims[layer], out_dim)
|
|
|
|
if weight_norm:
|
|
lin = nn.utils.weight_norm(lin)
|
|
|
|
setattr(self, 'lin' + str(layer), lin)
|
|
|
|
self.relu = nn.ReLU()
|
|
|
|
def forward(self, points, normals, view_dirs, feature_vectors):
|
|
if self.embedview_fn is not None:
|
|
view_dirs = self.embedview_fn(view_dirs)
|
|
|
|
rendering_input = None
|
|
|
|
if self.mode == 'idr':
|
|
rendering_input = torch.cat(
|
|
[points, view_dirs, normals, feature_vectors], dim=-1)
|
|
elif self.mode == 'no_view_dir':
|
|
rendering_input = torch.cat([points, normals, feature_vectors],
|
|
dim=-1)
|
|
elif self.mode == 'no_normal':
|
|
rendering_input = torch.cat([points, view_dirs, feature_vectors],
|
|
dim=-1)
|
|
|
|
x = rendering_input
|
|
|
|
for layer in range(0, self.num_layers - 1):
|
|
lin = getattr(self, 'lin' + str(layer))
|
|
x = lin(x)
|
|
if layer < self.num_layers - 2:
|
|
x = self.relu(x)
|
|
if self.squeeze_out:
|
|
x = torch.sigmoid(x)
|
|
return x
|
|
|
|
|
|
class NeRF(nn.Module):
|
|
|
|
def __init__(self,
|
|
D=8,
|
|
W=256,
|
|
d_in=3,
|
|
d_in_view=3,
|
|
multires=0,
|
|
multires_view=0,
|
|
output_ch=4,
|
|
skips=[4],
|
|
use_viewdirs=False):
|
|
super(NeRF, self).__init__()
|
|
self.D = D
|
|
self.W = W
|
|
self.d_in = d_in
|
|
self.d_in_view = d_in_view
|
|
self.input_ch = 3
|
|
self.input_ch_view = 3
|
|
self.embed_fn = None
|
|
self.embed_fn_view = None
|
|
|
|
if multires > 0:
|
|
embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
|
|
self.embed_fn = embed_fn
|
|
self.input_ch = input_ch
|
|
|
|
if multires_view > 0:
|
|
embed_fn_view, input_ch_view = get_embedder(
|
|
multires_view, input_dims=d_in_view)
|
|
self.embed_fn_view = embed_fn_view
|
|
self.input_ch_view = input_ch_view
|
|
|
|
self.skips = skips
|
|
self.use_viewdirs = use_viewdirs
|
|
|
|
self.pts_linears = nn.ModuleList([nn.Linear(self.input_ch, W)] + [
|
|
nn.Linear(W, W) if i not in
|
|
self.skips else nn.Linear(W + self.input_ch, W)
|
|
for i in range(D - 1)
|
|
])
|
|
|
|
self.views_linears = nn.ModuleList(
|
|
[nn.Linear(self.input_ch_view + W, W // 2)])
|
|
|
|
if use_viewdirs:
|
|
self.feature_linear = nn.Linear(W, W)
|
|
self.alpha_linear = nn.Linear(W, 1)
|
|
self.rgb_linear = nn.Linear(W // 2, 3)
|
|
else:
|
|
self.output_linear = nn.Linear(W, output_ch)
|
|
|
|
def forward(self, input_pts, input_views):
|
|
if self.embed_fn is not None:
|
|
input_pts = self.embed_fn(input_pts)
|
|
if self.embed_fn_view is not None:
|
|
input_views = self.embed_fn_view(input_views)
|
|
|
|
h = input_pts
|
|
for i, l in enumerate(self.pts_linears):
|
|
h = self.pts_linears[i](h)
|
|
h = F.relu(h)
|
|
if i in self.skips:
|
|
h = torch.cat([input_pts, h], -1)
|
|
|
|
if self.use_viewdirs:
|
|
alpha = self.alpha_linear(h)
|
|
feature = self.feature_linear(h)
|
|
h = torch.cat([feature, input_views], -1)
|
|
|
|
for i, l in enumerate(self.views_linears):
|
|
h = self.views_linears[i](h)
|
|
h = F.relu(h)
|
|
|
|
rgb = self.rgb_linear(h)
|
|
return alpha, rgb
|
|
else:
|
|
assert False
|
|
|
|
|
|
class SingleVarianceNetwork(nn.Module):
|
|
|
|
def __init__(self, init_val):
|
|
super(SingleVarianceNetwork, self).__init__()
|
|
self.register_parameter('variance',
|
|
nn.Parameter(torch.tensor(init_val)))
|
|
|
|
def forward(self, x):
|
|
return torch.ones([len(x), 1],
|
|
device=self.variance.device) * torch.exp(
|
|
self.variance * 10.0)
|
|
|
|
|
|
class Mean(nn.Module):
|
|
|
|
def __init__(self, dim: list, keepdim=False):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.keepdim = keepdim
|
|
|
|
def forward(self, x):
|
|
return torch.mean(x, self.dim, self.keepdim)
|
|
|
|
|
|
class Discriminator(nn.Module):
|
|
|
|
def __init__(self, channel=32, patch=True):
|
|
super().__init__()
|
|
self.imsize = 32
|
|
self.nc = 3
|
|
|
|
self.channel = channel
|
|
self.patch = patch
|
|
in_channel = 3
|
|
layer = []
|
|
for idx in range(3):
|
|
layer.extend([
|
|
spectral_norm(
|
|
nn.Conv2d(
|
|
in_channel, channel * (2**idx), 3, stride=2,
|
|
padding=1)),
|
|
nn.LeakyReLU(inplace=True),
|
|
spectral_norm(
|
|
nn.Conv2d(
|
|
channel * (2**idx),
|
|
channel * (2**idx),
|
|
3,
|
|
stride=1,
|
|
padding=1)),
|
|
nn.LeakyReLU(inplace=True),
|
|
])
|
|
in_channel = channel * (2**idx)
|
|
self.body = nn.Sequential(*layer)
|
|
if self.patch:
|
|
self.head = spectral_norm(nn.Conv2d(in_channel, 1, 1, padding=0))
|
|
else:
|
|
self.head = nn.Sequential(Mean([1, 2]), nn.Linear(in_channel, 1))
|
|
|
|
def forward(self, x):
|
|
x = x[:, :self.nc]
|
|
x = x.view(-1, self.imsize, self.imsize, self.nc).permute(0, 3, 1, 2)
|
|
x = self.body(x)
|
|
x = self.head(x)
|
|
return x
|
|
|
|
|
|
class Embedder:
|
|
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
self.create_embedding_fn()
|
|
|
|
def create_embedding_fn(self):
|
|
embed_fns = []
|
|
d = self.kwargs['input_dims']
|
|
out_dim = 0
|
|
if self.kwargs['include_input']:
|
|
embed_fns.append(lambda x: x)
|
|
out_dim += d
|
|
|
|
max_freq = self.kwargs['max_freq_log2']
|
|
N_freqs = self.kwargs['num_freqs']
|
|
|
|
if self.kwargs['log_sampling']:
|
|
freq_bands = 2.**torch.linspace(0., max_freq, N_freqs)
|
|
else:
|
|
freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)
|
|
|
|
for freq in freq_bands:
|
|
for p_fn in self.kwargs['periodic_fns']:
|
|
embed_fns.append(
|
|
lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
|
|
out_dim += d
|
|
|
|
self.embed_fns = embed_fns
|
|
self.out_dim = out_dim
|
|
|
|
def embed(self, inputs):
|
|
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
|
|
|
|
|
|
def get_embedder(multires, input_dims=3):
|
|
embed_kwargs = {
|
|
'include_input': True,
|
|
'input_dims': input_dims,
|
|
'max_freq_log2': multires - 1,
|
|
'num_freqs': multires,
|
|
'log_sampling': True,
|
|
'periodic_fns': [torch.sin, torch.cos],
|
|
}
|
|
|
|
embedder_obj = Embedder(**embed_kwargs)
|
|
|
|
def embed(x, eo=embedder_obj):
|
|
return eo.embed(x)
|
|
|
|
return embed, embedder_obj.out_dim
|