| 96 | |
| 97 | # Positional encoding (section 5.1 of NERF) |
| 98 | class Embedder: |
| 99 | def __init__(self, **kwargs): |
| 100 | self.kwargs = kwargs |
| 101 | self.N_freqs = 0 |
| 102 | self.N = -1 # epoch to max frequency, for Nerfie embedding only |
| 103 | self.create_embedding_fn() |
| 104 | |
| 105 | def create_embedding_fn(self): |
| 106 | embed_fns = [] |
| 107 | d = self.kwargs['input_dims'] |
| 108 | out_dim = 0 |
| 109 | if self.kwargs['include_input']: |
| 110 | embed_fns.append(lambda x : x) |
| 111 | out_dim += d |
| 112 | |
| 113 | max_freq = self.kwargs['max_freq_log2'] |
| 114 | self.N_freqs = self.kwargs['num_freqs'] |
| 115 | |
| 116 | if self.kwargs['log_sampling']: |
| 117 | freq_bands = 2.**torch.linspace(0., max_freq, steps=self.N_freqs) # tensor([ 1., 2., 4., 8., 16., 32., 64., 128., 256., 512.]) |
| 118 | else: |
| 119 | freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=self.N_freqs) |
| 120 | |
| 121 | for freq in freq_bands: # 10 iters for 3D location, 4 iters for 2D direction |
| 122 | for p_fn in self.kwargs['periodic_fns']: |
| 123 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) |
| 124 | out_dim += d |
| 125 | self.embed_fns = embed_fns |
| 126 | self.out_dim = out_dim |
| 127 | |
| 128 | def embed(self, inputs): |
| 129 | if self.kwargs['max_freq_log2'] != 0: |
| 130 | ret = torch.cat([fn(inputs) for fn in self.embed_fns], -1) # cos, sin embedding |
| 131 | else: |
| 132 | ret = inputs |
| 133 | return ret |
| 134 | |
| 135 | def get_embed_weight(self, epoch, num_freqs, N): |
| 136 | ''' Nerfie Paper Eq.(8) ''' |
| 137 | alpha = num_freqs * epoch / N |
| 138 | W_j = [] |
| 139 | for i in range(num_freqs): |
| 140 | tmp = torch.clamp(torch.Tensor([alpha - i]), 0, 1) |
| 141 | tmp2 = (1 - torch.cos(torch.Tensor([np.pi]) * tmp)) / 2 |
| 142 | W_j.append(tmp2) |
| 143 | return W_j |
| 144 | |
| 145 | def embed_DNeRF(self, inputs, epoch): |
| 146 | ''' Nerfie paper section 3.5 Coarse-to-Fine Deformation Regularization ''' |
| 147 | # get weight for each frequency band j |
| 148 | W_j = self.get_embed_weight(epoch, self.N_freqs, self.N) # W_j: [W_0, W_1, W_2, ..., W_{m-1}] |
| 149 | |
| 150 | # Fourier embedding |
| 151 | out = [] |
| 152 | for fn in self.embed_fns: # 17, embed_fns:[input, cos, sin, cos, sin, ..., cos, sin] |
| 153 | out.append(fn(inputs)) |
| 154 | |
| 155 | # apply weighted positional encoding, only to cos&sins |