-
Notifications
You must be signed in to change notification settings - Fork 103
/
nerf_network.py
142 lines (117 loc) · 4.89 KB
/
nerf_network.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import torch
import torch.nn as nn
# import torch.nn.functional as F
# import numpy as np
from collections import OrderedDict
import logging
logger = logging.getLogger(__package__)
class Embedder(nn.Module):
def __init__(self, input_dim, max_freq_log2, N_freqs,
log_sampling=True, include_input=True,
periodic_fns=(torch.sin, torch.cos)):
'''
:param input_dim: dimension of input to be embedded
:param max_freq_log2: log2 of max freq; min freq is 1 by default
:param N_freqs: number of frequency bands
:param log_sampling: if True, frequency bands are linerly sampled in log-space
:param include_input: if True, raw input is included in the embedding
:param periodic_fns: periodic functions used to embed input
'''
super().__init__()
self.input_dim = input_dim
self.include_input = include_input
self.periodic_fns = periodic_fns
self.out_dim = 0
if self.include_input:
self.out_dim += self.input_dim
self.out_dim += self.input_dim * N_freqs * len(self.periodic_fns)
if log_sampling:
self.freq_bands = 2. ** torch.linspace(0., max_freq_log2, N_freqs)
else:
self.freq_bands = torch.linspace(2. ** 0., 2. ** max_freq_log2, N_freqs)
self.freq_bands = self.freq_bands.numpy().tolist()
def forward(self, input):
'''
:param input: tensor of shape [..., self.input_dim]
:return: tensor of shape [..., self.out_dim]
'''
assert (input.shape[-1] == self.input_dim)
out = []
if self.include_input:
out.append(input)
for i in range(len(self.freq_bands)):
freq = self.freq_bands[i]
for p_fn in self.periodic_fns:
out.append(p_fn(input * freq))
out = torch.cat(out, dim=-1)
assert (out.shape[-1] == self.out_dim)
return out
# default tensorflow initialization of linear layers
def weights_init(m):
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight.data)
if m.bias is not None:
nn.init.zeros_(m.bias.data)
class MLPNet(nn.Module):
def __init__(self, D=8, W=256, input_ch=3, input_ch_viewdirs=3,
skips=[4], use_viewdirs=False):
'''
:param D: network depth
:param W: network width
:param input_ch: input channels for encodings of (x, y, z)
:param input_ch_viewdirs: input channels for encodings of view directions
:param skips: skip connection in network
:param use_viewdirs: if True, will use the view directions as input
'''
super().__init__()
self.input_ch = input_ch
self.input_ch_viewdirs = input_ch_viewdirs
self.use_viewdirs = use_viewdirs
self.skips = skips
self.base_layers = []
dim = self.input_ch
for i in range(D):
self.base_layers.append(
nn.Sequential(nn.Linear(dim, W), nn.ReLU())
)
dim = W
if i in self.skips and i != (D-1): # skip connection after i^th layer
dim += input_ch
self.base_layers = nn.ModuleList(self.base_layers)
# self.base_layers.apply(weights_init) # xavier init
sigma_layers = [nn.Linear(dim, 1), ] # sigma must be positive
self.sigma_layers = nn.Sequential(*sigma_layers)
# self.sigma_layers.apply(weights_init) # xavier init
# rgb color
rgb_layers = []
base_remap_layers = [nn.Linear(dim, 256), ]
self.base_remap_layers = nn.Sequential(*base_remap_layers)
# self.base_remap_layers.apply(weights_init)
dim = 256 + self.input_ch_viewdirs
for i in range(1):
rgb_layers.append(nn.Linear(dim, W // 2))
rgb_layers.append(nn.ReLU())
dim = W // 2
rgb_layers.append(nn.Linear(dim, 3))
rgb_layers.append(nn.Sigmoid()) # rgb values are normalized to [0, 1]
self.rgb_layers = nn.Sequential(*rgb_layers)
# self.rgb_layers.apply(weights_init)
def forward(self, input):
'''
:param input: [..., input_ch+input_ch_viewdirs]
:return [..., 4]
'''
input_pts = input[..., :self.input_ch]
base = self.base_layers[0](input_pts)
for i in range(len(self.base_layers)-1):
if i in self.skips:
base = torch.cat((input_pts, base), dim=-1)
base = self.base_layers[i+1](base)
sigma = self.sigma_layers(base)
sigma = torch.abs(sigma)
base_remap = self.base_remap_layers(base)
input_viewdirs = input[..., -self.input_ch_viewdirs:]
rgb = self.rgb_layers(torch.cat((base_remap, input_viewdirs), dim=-1))
ret = OrderedDict([('rgb', rgb),
('sigma', sigma.squeeze(-1))])
return ret