| 62 | |
| 63 | # Positional encoding (section 5.1) |
| 64 | class Embedder: |
| 65 | def __init__(self, **kwargs): |
| 66 | self.kwargs = kwargs |
| 67 | self.N_freqs = 0 |
| 68 | self.N = -1 # epoch to max frequency, for Nerfie embedding only |
| 69 | self.create_embedding_fn() |
| 70 | |
| 71 | def create_embedding_fn(self): |
| 72 | embed_fns = [] |
| 73 | d = self.kwargs['input_dims'] |
| 74 | out_dim = 0 |
| 75 | if self.kwargs['include_input']: |
| 76 | embed_fns.append(lambda x : x) |
| 77 | out_dim += d |
| 78 | |
| 79 | max_freq = self.kwargs['max_freq_log2'] |
| 80 | self.N_freqs = self.kwargs['num_freqs'] |
| 81 | |
| 82 | if self.kwargs['log_sampling']: |
| 83 | freq_bands = 2.**torch.linspace(0., max_freq, steps=self.N_freqs) # tensor([ 1., 2., 4., 8., 16., 32., 64., 128., 256., 512.]) |
| 84 | else: |
| 85 | freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=self.N_freqs) |
| 86 | |
| 87 | for freq in freq_bands: # 10 iters for 3D location, 4 iters for 2D direction |
| 88 | for p_fn in self.kwargs['periodic_fns']: |
| 89 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) |
| 90 | out_dim += d |
| 91 | self.embed_fns = embed_fns |
| 92 | self.out_dim = out_dim |
| 93 | |
| 94 | def embed(self, inputs): |
| 95 | # inputs [65536, 3] |
| 96 | if self.kwargs['max_freq_log2'] != 0: |
| 97 | ret = torch.cat([fn(inputs) for fn in self.embed_fns], -1) # cos, sin embedding # ret.shape [65536, 63] |
| 98 | else: |
| 99 | ret = inputs |
| 100 | return ret |
| 101 | |
| 102 | def get_embed_weight(self, epoch, num_freqs, N): |
| 103 | ''' Nerfie Paper Eq.(8) ''' |
| 104 | alpha = num_freqs * epoch / N |
| 105 | W_j = [] |
| 106 | for i in range(num_freqs): |
| 107 | tmp = torch.clamp(torch.Tensor([alpha - i]), 0, 1) |
| 108 | tmp2 = (1 - torch.cos(torch.Tensor([np.pi]) * tmp)) / 2 |
| 109 | W_j.append(tmp2) |
| 110 | return W_j |
| 111 | |
| 112 | def embed_DNeRF(self, inputs, epoch): |
| 113 | ''' Nerfie paper section 3.5 Coarse-to-Fine Deformation Regularization ''' |
| 114 | # get weight for each frequency band j |
| 115 | W_j = self.get_embed_weight(epoch, self.N_freqs, self.N) # W_j: [W_0, W_1, W_2, ..., W_{m-1}] |
| 116 | |
| 117 | # Fourier embedding |
| 118 | out = [] |
| 119 | for fn in self.embed_fns: # 17, embed_fns:[input, cos, sin, cos, sin, ..., cos, sin] |
| 120 | out.append(fn(inputs)) |
| 121 | |