import os
from typing import Optional,Tuple,List,Union,Callable

import numpy as np
import torch
from torch import nn
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from tqdm import trange

# 設(shè)置GPU還是CPU設(shè)備
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

data = np.load('tiny_nerf_data.npz') # 加載數(shù)據(jù)集
images = data['images']  # 圖像數(shù)據(jù)
poses = data['poses']  # 位姿數(shù)據(jù)
focal = data['focal']  # 焦距數(shù)值

print(f'Images shape: {images.shape}')
print(f'Poses shape: {poses.shape}')
print(f'Focal length: {focal}')

height, width = images.shape[1:3]
near, far = 2., 6.

n_training = 100 # 訓練數(shù)據(jù)數(shù)量
testimg_idx = 101 # 測試數(shù)據(jù)下標
testimg, testpose = images[testimg_idx], poses[testimg_idx]

plt.imshow(testimg)
print('Pose')
print(testpose)

# 方向數(shù)據(jù)
dirs = np.stack([np.sum([0, 0, -1] * pose[:3, :3], axis=-1) for pose in poses])
# 原點數(shù)據(jù)
origins = poses[:, :3, -1]

# 繪圖的設(shè)置

ax = plt.figure(figsize=(12, 8)).add_subplot(projectinotallow='3d')
_ = ax.quiver(
  origins[..., 0].flatten(),
  origins[..., 1].flatten(),
  origins[..., 2].flatten(),
  dirs[..., 0].flatten(),
  dirs[..., 1].flatten(),
  dirs[..., 2].flatten(), length=0.5, normalize=True)

ax.set_xlabel('X')

ax.set_ylabel('Y')

ax.set_zlabel('z')

plt.show()

def get_rays(
  height: int, # 圖像高度
  width: int, # 圖像寬帶
  focal_length: float, # 焦距
  c2w: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
  """
  通過每個像素和相機原點，找到射線的原點和方向。
 
 """


  # 應用針孔相機模型收集每個像素的方向
  i, j = torch.meshgrid(
      torch.arange(width, dtype=torch.float32).to(c2w),
      torch.arange(height, dtype=torch.float32).to(c2w),
      indexing='ij')
      
  i, j = i.transpose(-1, -2), j.transpose(-1, -2)

  # 方向數(shù)據(jù)
  directions = torch.stack([(i - width * .5) / focal_length,
                            -(j - height * .5) / focal_length,
                            -torch.ones_like(i)
                           ], dim=-1)



  # 用相機位姿求出方向
  rays_d = torch.sum(directions[..., None, :] * c2w[:3, :3], dim=-1)


  # 默認所有射線原點相同
  rays_o = c2w[:3, -1].expand(rays_d.shape)
  return rays_o, rays_d

# 轉(zhuǎn)為PyTorch的tensor
images = torch.from_numpy(data['images'][:n_training]).to(device)
poses = torch.from_numpy(data['poses']).to(device)
focal = torch.from_numpy(data['focal']).to(device)
testimg = torch.from_numpy(data['images'][testimg_idx]).to(device)
testpose = torch.from_numpy(data['poses'][testimg_idx]).to(device)


# 針對每個圖像獲取射線
height, width = images.shape[1:3]

with torch.no_grad():
  ray_origin, ray_direction = get_rays(height, width, focal, testpose)


print('Ray Origin')

print(ray_origin.shape)
print(ray_origin[height // 2, width // 2, :])
print('')


print('Ray Direction')

print(ray_direction.shape)
print(ray_direction[height // 2, width // 2, :])
print('')

# 采樣函數(shù)定義
def sample_stratified(
  rays_o: torch.Tensor, # 射線原點
  rays_d: torch.Tensor, # 射線方向
  near: float,
  far: float,
  n_samples: int, # 采樣數(shù)量
  perturb: Optional[bool] = True, # 擾動設(shè)置
  inverse_depth: bool = False  # 反向深度
) -> Tuple[torch.Tensor, torch.Tensor]:
  """
  從規(guī)則的bin中沿著射線進行采樣。
  """

  # 沿著射線抓取采樣點
  t_vals = torch.linspace(0., 1., n_samples, device=rays_o.device)
  if not inverse_depth:
    # 由遠到近線性采樣
    z_vals = near * (1.-t_vals) + far * (t_vals)
  else:
    # 在反向深度中線性采樣
  z_vals = 1./(1./near * (1.-t_vals) + 1./far * (t_vals))

  # 沿著射線從bins中統(tǒng)一采樣
  if perturb:
    mids = .5 * (z_vals[1:] + z_vals[:-1])
    upper = torch.concat([mids, z_vals[-1:]], dim=-1)
    lower = torch.concat([z_vals[:1], mids], dim=-1)
    t_rand = torch.rand([n_samples], device=z_vals.device)
    z_vals = lower + (upper - lower) * t_rand
  z_vals = z_vals.expand(list(rays_o.shape[:-1]) + [n_samples])

  # 應用相應的縮放參數(shù)
  pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals[..., :, None]
  return pts, z_vals

y_vals = torch.zeros_like(z_vals)
# 調(diào)用采樣策略函數(shù)
_, z_vals_unperturbed = sample_stratified(rays_o, rays_d, near, far, n_samples,
                                  perturb=False, inverse_depth=inverse_depth)
# 繪圖相關(guān)plt.plot(z_vals_unperturbed[0].cpu().numpy(), 1 + y_vals[0].cpu().numpy(), 'b-o')
plt.plot(z_vals[0].cpu().numpy(), y_vals[0].cpu().numpy(), 'r-o')
plt.ylim([-1, 2])
plt.title('Stratified Sampling (blue) with Perturbation (red)')
ax = plt.gca()
ax.axes.yaxis.set_visible(False)
plt.grid(True)

# 位置編碼類
class PositionalEncoder(nn.Module):
  """
  對輸入點，做sine或者consine位置編碼。
  """
  def __init__(
    self,
    d_input: int,
    n_freqs: int,
    log_space: bool = False
  ):
    super().__init__()
    self.d_input = d_input
    self.n_freqs = n_freqs
    self.log_space = log_space
    self.d_output = d_input * (1 + 2 * self.n_freqs)
    self.embed_fns = [lambda x: x]

    # 定義線性或者log尺度的頻率
    if self.log_space:
      freq_bands = 2.**torch.linspace(0., self.n_freqs - 1, self.n_freqs)
    else:
      freq_bands = torch.linspace(2.**0., 2.**(self.n_freqs - 1), self.n_freqs)

    # 替換sin和cos
    for freq in freq_bands:
      self.embed_fns.append(lambda x, freq=freq: torch.sin(x * freq))
      self.embed_fns.append(lambda x, freq=freq: torch.cos(x * freq))
  
  def forward(
    self,
    x
  ) -> torch.Tensor:
    """
    實際使用位置編碼的函數(shù)。
    """
    return torch.concat([fn(x) for fn in self.embed_fns], dim=-1)

# 定義NeRF模型
class NeRF(nn.Module):
  """
  神經(jīng)輻射場模塊。
  """
  def __init__(
    self,
    d_input: int = 3,
    n_layers: int = 8,
    d_filter: int = 256,
    skip: Tuple[int] = (4,),
    d_viewdirs: Optional[int] = None
  ):
    super().__init__()
    self.d_input = d_input # 輸入
    self.skip = skip # 殘差連接
    self.act = nn.functional.relu # 激活函數(shù)
    self.d_viewdirs = d_viewdirs # 視圖方向

    # 創(chuàng)建模型的層結(jié)構(gòu)
    self.layers = nn.ModuleList(
      [nn.Linear(self.d_input, d_filter)] +
      [nn.Linear(d_filter + self.d_input, d_filter) if i in skip \
       else nn.Linear(d_filter, d_filter) for i in range(n_layers - 1)]
    )

    # Bottleneck 層
    if self.d_viewdirs is not None:
      # 如果使用視圖方向，分離alpha和RGB
      self.alpha_out = nn.Linear(d_filter, 1)
      self.rgb_filters = nn.Linear(d_filter, d_filter)
      self.branch = nn.Linear(d_filter + self.d_viewdirs, d_filter // 2)
      self.output = nn.Linear(d_filter // 2, 3)
    else:
      # 如果不使用試圖方向，則簡單輸出
      self.output = nn.Linear(d_filter, 4)
  
  def forward(
    self,
    x: torch.Tensor,
    viewdirs: Optional[torch.Tensor] = None
  ) -> torch.Tensor:
    r"""
    帶有視圖方向的前向傳播
    """

    # 判斷是否設(shè)置視圖方向
    if self.d_viewdirs is None and viewdirs is not None:
      raise ValueError('Cannot input x_direction if d_viewdirs was not given.')

    # 運行bottleneck層之前的網(wǎng)絡層
    x_input = x
    for i, layer in enumerate(self.layers):
      x = self.act(layer(x))
      if i in self.skip:
        x = torch.cat([x, x_input], dim=-1)

    # 運行 bottleneck
    if self.d_viewdirs is not None:
      # Split alpha from network output
      alpha = self.alpha_out(x)

      # 結(jié)果傳入到rgb過濾器
      x = self.rgb_filters(x)
      x = torch.concat([x, viewdirs], dim=-1)
      x = self.act(self.branch(x))
      x = self.output(x)

      # 拼接alpha一起作為輸出
      x = torch.concat([x, alpha], dim=-1)
    else:
      # 不拼接，簡單輸出
      x = self.output(x)
    return x

# 體積渲染
def cumprod_exclusive(
  tensor: torch.Tensor
) -> torch.Tensor:
  """
  (Courtesy of https://github.com/krrish94/nerf-pytorch)

  和tf.math.cumprod(..., exclusive=True)功能類似

  參數(shù):
  tensor (torch.Tensor): Tensor whose cumprod (cumulative product, see `torch.cumprod`) along dim=-1
    is to be computed.
  返回值:
  cumprod (torch.Tensor): cumprod of Tensor along dim=-1, mimiciking the functionality of
    tf.math.cumprod(..., exclusive=True) (see `tf.math.cumprod` for details).
  """

  # 首先計算規(guī)則的cunprod
  cumprod = torch.cumprod(tensor, -1)
  cumprod = torch.roll(cumprod, 1, -1)
  # 用1替換首個元素
  cumprod[..., 0] = 1.
  
  return cumprod

# 輸出到圖像的函數(shù)
def raw2outputs(
  raw: torch.Tensor,
  z_vals: torch.Tensor,
  rays_d: torch.Tensor,
  raw_noise_std: float = 0.0,
  white_bkgd: bool = False
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
  """
  將NeRF的輸出轉(zhuǎn)換為RGB輸出。
  """

  # 沿著`z_vals`軸元素之間的差值.
  dists = z_vals[..., 1:] - z_vals[..., :-1]
  dists = torch.cat([dists, 1e10 * torch.ones_like(dists[..., :1])], dim=-1)

  # 將每個距離乘以相應方向射線的法線，轉(zhuǎn)換為現(xiàn)實世界中的距離（考慮非單位方向）。
  dists = dists * torch.norm(rays_d[..., None, :], dim=-1)

  # 為模型預測密度添加噪音?？捎糜谠谟柧氝^程中對網(wǎng)絡進行正則化（防止出現(xiàn)浮點偽影）。
  noise = 0.
  if raw_noise_std > 0.:
    noise = torch.randn(raw[..., 3].shape) * raw_noise_std

  # Predict density of each sample along each ray. Higher values imply
  # higher likelihood of being absorbed at this point. [n_rays, n_samples]
  alpha = 1.0 - torch.exp(-nn.functional.relu(raw[..., 3] + noise) * dists)

  # 預測每條射線上每個樣本的密度。數(shù)值越大，表示該點被吸收的可能性越大。[n_ 射線，n_樣本］
  weights = alpha * cumprod_exclusive(1. - alpha + 1e-10)

  # 計算RGB圖的權(quán)重。
  rgb = torch.sigmoid(raw[..., :3])  # [n_rays, n_samples, 3]
  rgb_map = torch.sum(weights[..., None] * rgb, dim=-2)  # [n_rays, 3]

  # 估計預測距離的深度圖。
  depth_map = torch.sum(weights * z_vals, dim=-1)

  # 稀疏圖
  disp_map = 1. / torch.max(1e-10 * torch.ones_like(depth_map),
                            depth_map / torch.sum(weights, -1))

  # 沿著每條射線加權(quán)。
  acc_map = torch.sum(weights, dim=-1)

  # 要合成到白色背景上，請使用累積的 alpha 貼圖。
  if white_bkgd:
    rgb_map = rgb_map + (1. - acc_map[..., None])

  return rgb_map, depth_map, acc_map, weights

# 采樣概率密度函數(shù)
def sample_pdf(
  bins: torch.Tensor,
  weights: torch.Tensor,
  n_samples: int,
  perturb: bool = False
) -> torch.Tensor:
  """
  應用反向轉(zhuǎn)換采樣到一組加權(quán)點。
  """

  # 正則化權(quán)重得到概率密度函數(shù)。
  pdf = (weights + 1e-5) / torch.sum(weights + 1e-5, -1, keepdims=True) # [n_rays, weights.shape[-1]]

  # 將概率密度函數(shù)轉(zhuǎn)為累計分布函數(shù)。
  cdf = torch.cumsum(pdf, dim=-1) # [n_rays, weights.shape[-1]]
  cdf = torch.concat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1) # [n_rays, weights.shape[-1] + 1]

  # 從累計分布函數(shù)中提取樣本位置。perturb == 0 時為線性。
  if not perturb:
    u = torch.linspace(0., 1., n_samples, device=cdf.device)
    u = u.expand(list(cdf.shape[:-1]) + [n_samples]) # [n_rays, n_samples]
  else:
    u = torch.rand(list(cdf.shape[:-1]) + [n_samples], device=cdf.device) # [n_rays, n_samples]

  # 沿累計分布函數(shù)找出 u 值所在的索引。
  u = u.contiguous() # 返回具有相同值的連續(xù)張量。
  inds = torch.searchsorted(cdf, u, right=True) # [n_rays, n_samples]

  # 夾住超出范圍的索引。
  below = torch.clamp(inds - 1, min=0)
  above = torch.clamp(inds, max=cdf.shape[-1] - 1)
  inds_g = torch.stack([below, above], dim=-1) # [n_rays, n_samples, 2]

  # 從累計分布函數(shù)和相應的 bin 中心取樣。
  matched_shape = list(inds_g.shape[:-1]) + [cdf.shape[-1]]
  cdf_g = torch.gather(cdf.unsqueeze(-2).expand(matched_shape), dim=-1,
                       index=inds_g)
  bins_g = torch.gather(bins.unsqueeze(-2).expand(matched_shape), dim=-1,
                        index=inds_g)

  # 將樣本轉(zhuǎn)換為射線長度。
  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])

  return samples # [n_rays, n_samples]

def get_chunks(
  inputs: torch.Tensor,
  chunksize: int = 2**15
) -> List[torch.Tensor]:
  """
  輸入分塊。
  """
  return [inputs[i:i + chunksize] for i in range(0, inputs.shape[0], chunksize)]

def prepare_chunks(
  points: torch.Tensor,
  encoding_function: Callable[[torch.Tensor], torch.Tensor],
  chunksize: int = 2**15
) -> List[torch.Tensor]:
  """
  對點進行編碼和分塊，為 NeRF 模型做好準備。
  """
  points = points.reshape((-1, 3))
  points = encoding_function(points)
  points = get_chunks(points, chunksize=chunksize)
  return points

def prepare_viewdirs_chunks(
  points: torch.Tensor,
  rays_d: torch.Tensor,
  encoding_function: Callable[[torch.Tensor], torch.Tensor],
  chunksize: int = 2**15
) -> List[torch.Tensor]:
  r"""
  對視圖方向進行編碼和分塊，為 NeRF 模型做好準備。
  """
  viewdirs = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)
  viewdirs = viewdirs[:, None, ...].expand(points.shape).reshape((-1, 3))
  viewdirs = encoding_function(viewdirs)
  viewdirs = get_chunks(viewdirs, chunksize=chunksize)
  return viewdirs
def nerf_forward(
  rays_o: torch.Tensor,
  rays_d: torch.Tensor,
  near: float,
  far: float,
  encoding_fn: Callable[[torch.Tensor], torch.Tensor],
  coarse_model: nn.Module,
  kwargs_sample_stratified: dict = None,
  n_samples_hierarchical: int = 0,
  kwargs_sample_hierarchical: dict = None,
  fine_model = None,
  viewdirs_encoding_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
  chunksize: int = 2**15
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict]:
  """
  計算一次前向傳播
  """

  # 設(shè)置參數(shù)
  if kwargs_sample_stratified is None:
    kwargs_sample_stratified = {}
  if kwargs_sample_hierarchical is None:
    kwargs_sample_hierarchical = {}
  
  # 沿著每條射線的樣本查詢點。
  query_points, z_vals = sample_stratified(
      rays_o, rays_d, near, far, **kwargs_sample_stratified)

  # 準備批次。
  batches = prepare_chunks(query_points, encoding_fn, chunksize=chunksize)
  if viewdirs_encoding_fn is not None:
    batches_viewdirs = prepare_viewdirs_chunks(query_points, rays_d,
                                               viewdirs_encoding_fn,
                                               chunksize=chunksize)
  else:
    batches_viewdirs = [None] * len(batches)

  # 稀疏模型流程。
  predictions = []
  for batch, batch_viewdirs in zip(batches, batches_viewdirs):
    predictions.append(coarse_model(batch, viewdirs=batch_viewdirs))
  raw = torch.cat(predictions, dim=0)
  raw = raw.reshape(list(query_points.shape[:2]) + [raw.shape[-1]])

  # 執(zhí)行可微分體積渲染，重新合成 RGB 圖像。
  rgb_map, depth_map, acc_map, weights = raw2outputs(raw, z_vals, rays_d)
  
  outputs = {
      'z_vals_stratified': z_vals
  }

  if n_samples_hierarchical > 0:
    # Save previous outputs to return.
    rgb_map_0, depth_map_0, acc_map_0 = rgb_map, depth_map, acc_map

    # 對精細查詢點進行分層抽樣。
    query_points, z_vals_combined, z_hierarch = sample_hierarchical(
      rays_o, rays_d, z_vals, weights, n_samples_hierarchical,
      **kwargs_sample_hierarchical)

    # 像以前一樣準備輸入。
    batches = prepare_chunks(query_points, encoding_fn, chunksize=chunksize)
    if viewdirs_encoding_fn is not None:
      batches_viewdirs = prepare_viewdirs_chunks(query_points, rays_d,
                                                 viewdirs_encoding_fn,
                                                 chunksize=chunksize)
    else:
      batches_viewdirs = [None] * len(batches)

    # 通過精細模型向前傳遞新樣本。
    fine_model = fine_model if fine_model is not None else coarse_model
    predictions = []
    for batch, batch_viewdirs in zip(batches, batches_viewdirs):
      predictions.append(fine_model(batch, viewdirs=batch_viewdirs))
    raw = torch.cat(predictions, dim=0)
    raw = raw.reshape(list(query_points.shape[:2]) + [raw.shape[-1]])

    # 執(zhí)行可微分體積渲染，重新合成 RGB 圖像。
    rgb_map, depth_map, acc_map, weights = raw2outputs(raw, z_vals_combined, rays_d)
    
    # 存儲輸出
    outputs['z_vals_hierarchical'] = z_hierarch
    outputs['rgb_map_0'] = rgb_map_0
    outputs['depth_map_0'] = depth_map_0
    outputs['acc_map_0'] = acc_map_0

  # 存儲輸出
  outputs['rgb_map'] = rgb_map
  outputs['depth_map'] = depth_map
  outputs['acc_map'] = acc_map
  outputs['weights'] = weights
  return outputs

# 編碼器
d_input = 3           # 輸入維度
n_freqs = 10          # 輸入到編碼函數(shù)中的樣本點數(shù)量
log_space = True      # 如果設(shè)置，頻率按對數(shù)空間縮放
use_viewdirs = True   # 如果設(shè)置，則使用視圖方向作為輸入
n_freqs_views = 4     # 視圖編碼功能的數(shù)量

# 采樣策略
n_samples = 64         # 每條射線的空間樣本數(shù)
perturb = True         # 如果設(shè)置，則對采樣位置應用噪聲
inverse_depth = False  # 如果設(shè)置，則按反深度線性采樣點

# 模型
d_filter = 128          # 線性層濾波器的尺寸
n_layers = 2            # bottleneck層數(shù)量
skip = []               # 應用輸入殘差的層級
use_fine_model = True   # 如果設(shè)置，則創(chuàng)建一個精細模型
d_filter_fine = 128     # 精細網(wǎng)絡線性層濾波器的尺寸
n_layers_fine = 6       # 精細網(wǎng)絡瓶頸層數(shù)

# 分層采樣
n_samples_hierarchical = 64   # 每條射線的樣本數(shù)
perturb_hierarchical = False  # 如果設(shè)置，則對采樣位置應用噪聲

# 優(yōu)化器
lr = 5e-4  # 學習率

# 訓練
n_iters = 10000
batch_size = 2**14          # 每個梯度步長的射線數(shù)量（2 的冪次）
one_image_per_step = True   # 每個梯度步驟一個圖像（禁用批處理）
chunksize = 2**14           # 根據(jù)需要進行修改，以適應 GPU 內(nèi)存
center_crop = True          # 裁剪圖像的中心部分（每幅圖像裁剪一次）
center_crop_iters = 50      # 經(jīng)過這么多epoch后，停止裁剪中心
display_rate = 25          # 每 X 個epoch顯示一次測試輸出

# 早停
warmup_iters = 100          # 熱身階段的迭代次數(shù)
warmup_min_fitness = 10.0   # 在熱身_iters 處繼續(xù)訓練的最小 PSNR 值
n_restarts = 10             # 訓練停滯時重新開始的次數(shù)

# 捆綁了各種函數(shù)的參數(shù)，以便一次性傳遞。
kwargs_sample_stratified = {
    'n_samples': n_samples,
    'perturb': perturb,
    'inverse_depth': inverse_depth
}
kwargs_sample_hierarchical = {
    'perturb': perturb
}

# 繪制采樣函數(shù)
def plot_samples(
  z_vals: torch.Tensor,
  z_hierarch: Optional[torch.Tensor] = None,
  ax: Optional[np.ndarray] = None):
  r"""
  繪制分層樣本和（可選）分級樣本。
  """
  y_vals = 1 + np.zeros_like(z_vals)

  if ax is None:
    ax = plt.subplot()
  ax.plot(z_vals, y_vals, 'b-o')
  if z_hierarch is not None:
    y_hierarch = np.zeros_like(z_hierarch)
    ax.plot(z_hierarch, y_hierarch, 'r-o')
  ax.set_ylim([-1, 2])
  ax.set_title('Stratified  Samples (blue) and Hierarchical Samples (red)')
  ax.axes.yaxis.set_visible(False)
  ax.grid(True)
  return ax

def crop_center(
  img: torch.Tensor,
  frac: float = 0.5
) -> torch.Tensor:
  r"""
  從圖像中裁剪中心方形。
  """
  h_offset = round(img.shape[0] * (frac / 2))
  w_offset = round(img.shape[1] * (frac / 2))
  return img[h_offset:-h_offset, w_offset:-w_offset]

class EarlyStopping:
  r"""
  基于適配標準的早期停止輔助器
  """
  def __init__(
    self,
    patience: int = 30,
    margin: float = 1e-4
  ):
    self.best_fitness = 0.0
    self.best_iter = 0
    self.margin = margin
    self.patience = patience or float('inf')  # 在epoch停止提高后等待的停止時間

  def __call__(
    self,
    iter: int,
    fitness: float
  ):
    r"""
    檢查是否符合停止標準。
    """
    if (fitness - self.best_fitness) > self.margin:
      self.best_iter = iter
      self.best_fitness = fitness
    delta = iter - self.best_iter
    stop = delta >= self.patience  # 超過耐性則停止訓練
    return stop
def init_models():
  r"""
  為 NeRF 訓練初始化模型、編碼器和優(yōu)化器。
  """
  # 編碼器
  encoder = PositionalEncoder(d_input, n_freqs, log_space=log_space)
  encode = lambda x: encoder(x)

  # 視圖方向編碼
  if use_viewdirs:
    encoder_viewdirs = PositionalEncoder(d_input, n_freqs_views,
                                        log_space=log_space)
    encode_viewdirs = lambda x: encoder_viewdirs(x)
    d_viewdirs = encoder_viewdirs.d_output
  else:
    encode_viewdirs = None
    d_viewdirs = None

  # 模型
  model = NeRF(encoder.d_output, n_layers=n_layers, d_filter=d_filter, skip=skip,
              d_viewdirs=d_viewdirs)
  model.to(device)
  model_params = list(model.parameters())
  if use_fine_model:
    fine_model = NeRF(encoder.d_output, n_layers=n_layers, d_filter=d_filter, skip=skip,
                      d_viewdirs=d_viewdirs)
    fine_model.to(device)
    model_params = model_params + list(fine_model.parameters())
  else:
    fine_model = None

  # 優(yōu)化器
  optimizer = torch.optim.Adam(model_params, lr=lr)

  # 早停
  warmup_stopper = EarlyStopping(patience=50)

  return model, fine_model, encode, encode_viewdirs, optimizer, warmup_stopper

def train():
  r"""
  啟動 NeRF 訓練。
  """
  # 對所有圖像進行射線洗牌。
  if not one_image_per_step:
    height, width = images.shape[1:3]
    all_rays = torch.stack([torch.stack(get_rays(height, width, focal, p), 0)
                        for p in poses[:n_training]], 0)
    rays_rgb = torch.cat([all_rays, images[:, None]], 1)
    rays_rgb = torch.permute(rays_rgb, [0, 2, 3, 1, 4])
    rays_rgb = rays_rgb.reshape([-1, 3, 3])
    rays_rgb = rays_rgb.type(torch.float32)
    rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
    i_batch = 0

  train_psnrs = []
  val_psnrs = []
  iternums = []
  for i in trange(n_iters):
    model.train()

    if one_image_per_step:
      # 隨機選擇一張圖片作為目標。
      target_img_idx = np.random.randint(images.shape[0])
      target_img = images[target_img_idx].to(device)
      if center_crop and i < center_crop_iters:
        target_img = crop_center(target_img)
      height, width = target_img.shape[:2]
      target_pose = poses[target_img_idx].to(device)
      rays_o, rays_d = get_rays(height, width, focal, target_pose)
      rays_o = rays_o.reshape([-1, 3])
      rays_d = rays_d.reshape([-1, 3])
    else:
      # 在所有圖像上隨機顯示。
      batch = rays_rgb[i_batch:i_batch + batch_size]
      batch = torch.transpose(batch, 0, 1)
      rays_o, rays_d, target_img = batch
      height, width = target_img.shape[:2]
      i_batch += batch_size
      # 一個epoch后洗牌
      if i_batch >= rays_rgb.shape[0]:
          rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
          i_batch = 0
    target_img = target_img.reshape([-1, 3])

    # 運行 TinyNeRF 的一次迭代，得到渲染后的 RGB 圖像。
    outputs = nerf_forward(rays_o, rays_d,
                           near, far, encode, model,
                           kwargs_sample_stratified=kwargs_sample_stratified,
                           n_samples_hierarchical=n_samples_hierarchical,
                           kwargs_sample_hierarchical=kwargs_sample_hierarchical,
                           fine_model=fine_model,
                           viewdirs_encoding_fn=encode_viewdirs,
                           chunksize=chunksize)
    
    # 檢查任何數(shù)字問題。
    for k, v in outputs.items():
      if torch.isnan(v).any():
        print(f"! [Numerical Alert] {k} contains NaN.")
      if torch.isinf(v).any():
        print(f"! [Numerical Alert] {k} contains Inf.")

    # 反向傳播
    rgb_predicted = outputs['rgb_map']
    loss = torch.nn.functional.mse_loss(rgb_predicted, target_img)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    psnr = -10. * torch.log10(loss)
    train_psnrs.append(psnr.item())

    # 以給定的顯示速率評估測試值。
    if i % display_rate == 0:
      model.eval()
      height, width = testimg.shape[:2]
      rays_o, rays_d = get_rays(height, width, focal, testpose)
      rays_o = rays_o.reshape([-1, 3])
      rays_d = rays_d.reshape([-1, 3])
      outputs = nerf_forward(rays_o, rays_d,
                             near, far, encode, model,
                             kwargs_sample_stratified=kwargs_sample_stratified,
                             n_samples_hierarchical=n_samples_hierarchical,
                             kwargs_sample_hierarchical=kwargs_sample_hierarchical,
                             fine_model=fine_model,
                             viewdirs_encoding_fn=encode_viewdirs,
                             chunksize=chunksize)

      rgb_predicted = outputs['rgb_map']
      loss = torch.nn.functional.mse_loss(rgb_predicted, testimg.reshape(-1, 3))
      print("Loss:", loss.item())
      val_psnr = -10. * torch.log10(loss)
      val_psnrs.append(val_psnr.item())
      iternums.append(i)

      # 繪制輸出示例
      fig, ax = plt.subplots(1, 4, figsize=(24,4), gridspec_kw={'width_ratios': [1, 1, 1, 3]})
      ax[0].imshow(rgb_predicted.reshape([height, width, 3]).detach().cpu().numpy())
      ax[0].set_title(f'Iteration: {i}')
      ax[1].imshow(testimg.detach().cpu().numpy())
      ax[1].set_title(f'Target')
      ax[2].plot(range(0, i + 1), train_psnrs, 'r')
      ax[2].plot(iternums, val_psnrs, 'b')
      ax[2].set_title('PSNR (train=red, val=blue')
      z_vals_strat = outputs['z_vals_stratified'].view((-1, n_samples))
      z_sample_strat = z_vals_strat[z_vals_strat.shape[0] // 2].detach().cpu().numpy()
      if 'z_vals_hierarchical' in outputs:
        z_vals_hierarch = outputs['z_vals_hierarchical'].view((-1, n_samples_hierarchical))
        z_sample_hierarch = z_vals_hierarch[z_vals_hierarch.shape[0] // 2].detach().cpu().numpy()
      else:
        z_sample_hierarch = None
      _ = plot_samples(z_sample_strat, z_sample_hierarch, ax=ax[3])
      ax[3].margins(0)
      plt.show()

    # 檢查 PSNR 是否存在問題，如果發(fā)現(xiàn)問題，則停止運行。
    if i == warmup_iters - 1:
      if val_psnr < warmup_min_fitness:
        print(f'Val PSNR {val_psnr} below warmup_min_fitness {warmup_min_fitness}. Stopping...')
        return False, train_psnrs, val_psnrs
    elif i < warmup_iters:
      if warmup_stopper is not None and warmup_stopper(i, psnr):
        print(f'Train PSNR flatlined at {psnr} for {warmup_stopper.patience} iters. Stopping...')
        return False, train_psnrs, val_psnrs
    
  return True, train_psnrs, val_psnrs