Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from torch import nn | |
| import numpy as np | |
| import math | |
| # import tinycudann as tcnn | |
| class SineLayer(nn.Module): | |
| # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of omega_0. | |
| # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the | |
| # nonlinearity. Different signals may require different omega_0 in the first layer - this is a | |
| # hyperparameter. | |
| # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of | |
| # activations constant, but boost gradients to the weight matrix (see supplement Sec. 1.5) | |
| def __init__(self, in_features, out_features, bias=True, | |
| is_first=False, omega_0=30): | |
| super().__init__() | |
| self.omega_0 = omega_0 | |
| self.is_first = is_first | |
| self.in_features = in_features | |
| self.linear = nn.Linear(in_features, out_features, bias=bias) | |
| self.init_weights() | |
| def init_weights(self): | |
| with torch.no_grad(): | |
| if self.is_first: | |
| self.linear.weight.uniform_(-1 / self.in_features, | |
| 1 / self.in_features) | |
| else: | |
| self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, | |
| np.sqrt(6 / self.in_features) / self.omega_0) | |
| def forward(self, input): | |
| return torch.sin(self.omega_0 * self.linear(input)) | |
| def forward_with_intermediate(self, input): | |
| # For visualization of activation distributions | |
| intermediate = self.omega_0 * self.linear(input) | |
| return torch.sin(intermediate), intermediate | |
| class Siren(nn.Module): | |
| def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=False, | |
| first_omega_0=30, hidden_omega_0=30.): | |
| super().__init__() | |
| self.net = [] | |
| self.net.append(SineLayer(in_features, hidden_features, | |
| is_first=True, omega_0=first_omega_0)) | |
| for i in range(hidden_layers): | |
| self.net.append(SineLayer(hidden_features, hidden_features, | |
| is_first=False, omega_0=hidden_omega_0)) | |
| if outermost_linear: | |
| final_linear = nn.Linear(hidden_features, out_features) | |
| with torch.no_grad(): | |
| final_linear.weight.uniform_(-np.sqrt(6 / hidden_features) / hidden_omega_0, | |
| np.sqrt(6 / hidden_features) / hidden_omega_0) | |
| self.net.append(final_linear) | |
| else: | |
| self.net.append(SineLayer(hidden_features, out_features, | |
| is_first=False, omega_0=hidden_omega_0)) | |
| self.net = nn.Sequential(*self.net) | |
| def forward(self, coords): | |
| output = self.net(coords) | |
| return output | |
| class Homography(nn.Module): | |
| def __init__(self, in_features=1, hidden_features=256, hidden_layers=1): | |
| super().__init__() | |
| out_features = 8 | |
| self.net = [] | |
| self.net.append(nn.Linear(in_features, hidden_features)) | |
| self.net.append(nn.ReLU(inplace=True)) | |
| for i in range(hidden_layers): | |
| self.net.append(nn.Linear(hidden_features, hidden_features)) | |
| self.net.append(nn.ReLU(inplace=True)) | |
| self.net.append(nn.Linear(hidden_features, out_features)) | |
| self.net = nn.Sequential(*self.net) | |
| self.init_weights() | |
| def init_weights(self): | |
| with torch.no_grad(): | |
| self.net[-1].bias.copy_(torch.Tensor([1., 0., 0., 0., 1., 0., 0., 0.])) | |
| def forward(self, coords): | |
| output = self.net(coords) | |
| return output | |
| class Annealed(nn.Module): | |
| def __init__(self, in_channels, annealed_step, annealed_begin_step=0, identity=True): | |
| """ | |
| Defines a function that embeds x to (x, sin(2^k x), cos(2^k x), ...) | |
| in_channels: number of input channels (3 for both xyz and direction) | |
| """ | |
| super(Annealed, self).__init__() | |
| self.N_freqs = 16 | |
| self.in_channels = in_channels | |
| self.annealed = True | |
| self.annealed_step = annealed_step | |
| self.annealed_begin_step = annealed_begin_step | |
| self.index = torch.linspace(0, self.N_freqs - 1, self.N_freqs) | |
| self.identity = identity | |
| self.index_2 = self.index.view(-1, 1).repeat(1, 2).view(-1) | |
| def forward(self, x_embed, step): | |
| """ | |
| Embeds x to (x, sin(2^k x), cos(2^k x), ...) | |
| Different from the paper, "x" is also in the output | |
| See https://github.com/bmild/nerf/issues/12 | |
| Inputs: | |
| x: (B, self.in_channels) | |
| Outputs: | |
| out: (B, self.out_channels) | |
| """ | |
| use_PE = False | |
| if self.annealed_begin_step == 0: | |
| # calculate the w for each freq bands | |
| alpha = self.N_freqs * step / float(self.annealed_step) | |
| else: | |
| if step <= self.annealed_begin_step: | |
| alpha = 0 | |
| else: | |
| alpha = (self.N_freqs) * (step - self.annealed_begin_step) / float( | |
| self.annealed_step) | |
| w = (1 - torch.cos(math.pi * torch.clamp(alpha * torch.ones_like(self.index_2) - self.index_2, 0, 1))) / 2 | |
| if use_PE: | |
| w[16:] = w[:16] | |
| out = x_embed * w.to(x_embed.device) | |
| return out | |
| class BARF_PE(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.encoder = tcnn.Encoding(n_input_dims=2, | |
| encoding_config=config["positional encoding"]) | |
| self.decoder = tcnn.Network(n_input_dims=self.encoder.n_output_dims + | |
| 2, | |
| n_output_dims=3, | |
| network_config=config["BARF network"]) | |
| def forward(self, x, step=0, aneal_func=None): | |
| input = x | |
| input = self.encoder(input) | |
| if aneal_func is not None: | |
| input = torch.cat([x, aneal_func(input,step)], dim=-1) | |
| else: | |
| input = torch.cat([x, input], dim=-1) | |
| weight = torch.ones(input.shape[-1], device=input.device).cuda() | |
| x = self.decoder(weight * input) | |
| return x | |
| class Deform_Hash3d(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.encoder = tcnn.Encoding(n_input_dims=3, | |
| encoding_config=config["encoding_deform3d"]) | |
| self.decoder = nn.Sequential(nn.Linear(self.encoder.n_output_dims + 3, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 256), | |
| nn.ReLU(), | |
| nn.Linear(256, 2) | |
| ) | |
| def forward(self, x, step=0, aneal_func=None): | |
| input = x | |
| input = self.encoder(input) | |
| if aneal_func is not None: | |
| input = torch.cat([x, aneal_func(input,step)], dim=-1) | |
| else: | |
| input = torch.cat([x, input], dim=-1) | |
| weight = torch.ones(input.shape[-1], device=input.device).cuda() | |
| x = self.decoder(weight * input) / 5 | |
| return x | |
| class Deform_Hash3d_Warp(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.Deform_Hash3d = Deform_Hash3d(config) | |
| def forward(self, xyt_norm, step=0,aneal_func=None): | |
| x = self.Deform_Hash3d(xyt_norm,step=step, aneal_func=aneal_func) | |
| return x |