# 构造位置编码转换器

def __init__(self, **kwargs):

# 传入参数

self.kwargs = kwargs

# 创建位置编码

self.create_embedding_fn()

def create_embedding_fn(self):

# 读取参数

embed_fns = []

d = self.kwargs['input_dims']

out_dim = 0

# 如果包含输入就将编码前的值并入位置编码后的结果

if self.kwargs['include_input']:

embed_fns.append(lambda x : x)

out_dim += d

max_freq = self.kwargs['max_freq_log2']

N_freqs = self.kwargs['num_freqs']

# 两种表达形式

if self.kwargs['log_sampling']:

freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs)

else:

freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs)

# 按照参考函数生成最后的位置编码，使用的是sin，cos。

# 并改变输出维度

for freq in freq_bands:

for p_fn in self.kwargs['periodic_fns']:

embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))

out_dim += d

self.embed_fns = embed_fns

self.out_dim = out_dim

def embed(self, inputs):

# 如果不使用位置编码返回三个通道，此时的位置编码为恒等函数，f(x) = x。

if i == -1:

return nn.Identity(), 3

# multiries 为频率，视角频率为 4，坐标为 10。

embed_kwargs = {

'include_input' : True,

'input_dims' : 3,

'max_freq_log2' : multires-1,

'num_freqs' : multires,

'log_sampling' : True,

# 可以更换位置编码的形式

'periodic_fns' : [torch.sin, torch.cos],

}

embedder_obj = Embedder(**embed_kwargs)

embed = lambda x, eo=embedder_obj : eo.embed(x)

# 定义 NeRF 网络

def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False):

"""

super(NeRF, self).__init__()

# 网络层数

self.D = D

# 网络宽度

self.W = W

# 输入维度

self.input_ch = input_ch

# 输入视角信息

self.input_ch_views = input_ch_views

# 拼接，论文的网络图里面每经过一定层数就会把输入值拼接到网络里

self.skips = skips

# 是否使用视角

self.use_viewdirs = use_viewdirs

# 第一层是 input_ch -> w

# 不拼接输入的层为 w -> w

# 需要拼接输入的层为 w + input_ch -> w

# 这里面写得比较简单，只使用了一个简单的列表推导式实现

self.pts_linears = nn.ModuleList(

[nn.Linear(input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + input_ch, W) for i in range(D-1)])

# 下面是输入视角信息的层，拼接了视角信息，输出减半

self.views_linears = nn.ModuleList([nn.Linear(input_ch_views + W, W//2)])

### Implementation according to the paper

# self.views_linears = nn.ModuleList(

# [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])

if use_viewdirs:

# 使用视角的话，先经过全连接层

self.feature_linear = nn.Linear(W, W)

# 分别输出体密度和 rgb

self.alpha_linear = nn.Linear(W, 1)

self.rgb_linear = nn.Linear(W//2, 3)

else:

# 否则输出一个四通道的张量，分别代表体密度和 RGB

self.output_linear = nn.Linear(W, output_ch)

# 前向传播

def forward(self, x):

# 将读入的位置信息和视角信息分开

input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1)

h = input_pts

# 先将位置信息经过全连接层

for i, l in enumerate(self.pts_linears):

h = self.pts_linears[i](h)

h = F.relu(h)

# 碰到 skip 就将输入并入

if i in self.skips:

h = torch.cat([input_pts, h], -1)

# 根据使用视角信息确定网络尾部

if self.use_viewdirs:

alpha = self.alpha_linear(h)

feature = self.feature_linear(h)

h = torch.cat([feature, input_views], -1)

for i, l in enumerate(self.views_linears):

h = self.views_linears[i](h)

h = F.relu(h)

rgb = self.rgb_linear(h)

outputs = torch.cat([rgb, alpha], -1)

else:

outputs = self.output_linear(h)

# 返回输出结果

return outputs

def load_weights_from_keras(self, weights):

assert self.use_viewdirs, "Not implemented if use_viewdirs=False"

# Load pts_linears

for i in range(self.D):

idx_pts_linears = 2 * i

self.pts_linears[i].weight.data = torch.from_numpy(np.transpose(weights[idx_pts_linears]))

self.pts_linears[i].bias.data = torch.from_numpy(np.transpose(weights[idx_pts_linears+1]))

# Load feature_linear

idx_feature_linear = 2 * self.D

self.feature_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_feature_linear]))

self.feature_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_feature_linear+1]))

# Load views_linears

idx_views_linears = 2 * self.D + 2

self.views_linears[0].weight.data = torch.from_numpy(np.transpose(weights[idx_views_linears]))

self.views_linears[0].bias.data = torch.from_numpy(np.transpose(weights[idx_views_linears+1]))

# Load rgb_linear

idx_rbg_linear = 2 * self.D + 4

self.rgb_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear]))

self.rgb_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear+1]))

# Load alpha_linear

idx_alpha_linear = 2 * self.D + 6

self.alpha_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear]))

i, j = torch.meshgrid(torch.linspace(0, W-1, W), torch.linspace(0, H-1, H)) # pytorch's meshgrid has indexing='ij'

i = i.t()

j = j.t()

dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -torch.ones_like(i)], -1)

# Rotate ray directions from camera frame to the world frame

rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) # dot product, equals to: [c2w.dot(dir) for dir in dirs]

# Translate camera frame's origin to the world frame. It is the origin of all rays.

rays_o = c2w[:3,-1].expand(rays_d.shape)

# 将每个像素的位置生成i、j

i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32), indexing='xy')

# 按照方向压缩坐标

dirs = np.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -np.ones_like(i)], -1)

# Rotate ray directions from camera frame to the world frame

# 将相机坐标系下的射线转换到世界坐标系下

rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) # dot product, equals to: [c2w.dot(dir) for dir in dirs]

# Translate camera frame's origin to the world frame. It is the origin of all rays.

# rays_o 为相机原点，每幅图像800*800对应一个相机原点

rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))

# Shift ray origins to near plane

t = -(near + rays_o[...,2]) / rays_d[...,2]

rays_o = rays_o + t[...,None] * rays_d

# Projection

o0 = -1./(W/(2.*focal)) * rays_o[...,0] / rays_o[...,2]

o1 = -1./(H/(2.*focal)) * rays_o[...,1] / rays_o[...,2]

o2 = 1. + 2. * near / rays_o[...,2]

d0 = -1./(W/(2.*focal)) * (rays_d[...,0]/rays_d[...,2] - rays_o[...,0]/rays_o[...,2])

d1 = -1./(H/(2.*focal)) * (rays_d[...,1]/rays_d[...,2] - rays_o[...,1]/rays_o[...,2])

d2 = -2. * near / rays_o[...,2]

rays_o = torch.stack([o0,o1,o2], -1)

rays_d = torch.stack([d0,d1,d2], -1)

# Get pdf

weights = weights + 1e-5 # prevent nans

pdf = weights / torch.sum(weights, -1, keepdim=True)

cdf = torch.cumsum(pdf, -1)

cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # (batch, len(bins))

# Take uniform samples

if det:

u = torch.linspace(0., 1., steps=N_samples)

u = u.expand(list(cdf.shape[:-1]) + [N_samples])

else:

u = torch.rand(list(cdf.shape[:-1]) + [N_samples])

# Pytest, overwrite u with numpy's fixed random numbers

if pytest:

np.random.seed(0)

new_shape = list(cdf.shape[:-1]) + [N_samples]

if det:

u = np.linspace(0., 1., N_samples)

u = np.broadcast_to(u, new_shape)

else:

u = np.random.rand(*new_shape)

u = torch.Tensor(u)

# Invert CDF

u = u.contiguous()

inds = torch.searchsorted(cdf, u, right=True)

below = torch.max(torch.zeros_like(inds-1), inds-1)

above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds)

inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)

# cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)

# bins_g = tf.gather(bins, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)

matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]

cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)

bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)

denom = (cdf_g[...,1]-cdf_g[...,0])

denom = torch.where(denom<1e-5, torch.ones_like(denom), denom)

t = (u-cdf_g[...,0])/denom

samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])