| 200 | # Positional encoding (section 5.1) |
| 201 | # borrow from nerf |
| 202 | class Embedder: |
| 203 | def __init__(self, **kwargs): |
| 204 | self.kwargs = kwargs |
| 205 | self.create_embedding_fn() |
| 206 | |
| 207 | def create_embedding_fn(self): |
| 208 | embed_fns = [] |
| 209 | d = self.kwargs['input_dims'] |
| 210 | out_dim = 0 |
| 211 | if self.kwargs['include_input']: |
| 212 | embed_fns.append(lambda x: x) |
| 213 | out_dim += d |
| 214 | |
| 215 | max_freq = self.kwargs['max_freq_log2'] |
| 216 | N_freqs = self.kwargs['num_freqs'] |
| 217 | |
| 218 | if self.kwargs['log_sampling']: |
| 219 | freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs) |
| 220 | else: |
| 221 | freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs) |
| 222 | |
| 223 | for freq in freq_bands: |
| 224 | for p_fn in self.kwargs['periodic_fns']: |
| 225 | embed_fns.append( |
| 226 | lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq)) |
| 227 | out_dim += d |
| 228 | |
| 229 | self.embed_fns = embed_fns |
| 230 | self.out_dim = out_dim |
| 231 | |
| 232 | def embed(self, inputs): |
| 233 | return torch.cat([fn(inputs) for fn in self.embed_fns], -1) |
| 234 | |
| 235 | |
| 236 | def get_embedder(multires, i=0): |