政安晨：【Keras机器学习示例演绎】（七）—

简介

设置

下载并加载数据

NeRF 模型

训练

可视化训练步骤

推理

渲染三维场景

可视化视频

结论

政安晨的个人主页：政安晨

欢迎 👍点赞✍评论⭐收藏

收录专栏: TensorFlow与Keras机器学习实战

希望政安晨的博客能够对您有所裨益，如有不足之处，欢迎在评论区提出指正！

本文目标：NeRF 中显示的体积渲染的最小实现。

简介

在本示例中，我们将介绍 Ben Mildenhall 等人的研究论文 NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis 的最小实现。

作者提出了一种巧妙的方法，即通过神经网络对体积场景函数进行建模，从而合成场景的新颖视图。

为了帮助大家直观地理解这个问题，让我们从下面这个问题开始：是否可以向神经网络提供图像中某个像素的位置，并要求网络预测该位置的颜色？

（下图：给定图像坐标的神经网络）

作为输入，并要求预测坐标处的颜色。

假设神经网络会记忆（过拟合）图像。这意味着，我们的神经网络将在其权重中编码整个图像。我们可以根据每个位置查询神经网络，它最终会重建整个图像。

（下图：训练有素的神经网络从头开始重现图像。）

现在的问题是，我们如何将这一想法扩展到三维体积场景的学习中？

要实现与上述类似的过程，就需要了解每个体素（体积像素）。

事实证明，这是一项相当具有挑战性的任务。

本文提出了一种使用少量场景图像学习三维场景的简约而优雅的方法。

他们放弃了使用体素进行训练。该网络学习对体积场景进行建模，从而生成三维场景的新视图（图像），而模型在训练时并未显示这些视图。

要充分了解这一过程，需要了解一些先决条件。我们将以这样一种方式来安排示例，以便您在开始实施前掌握所有必要的知识。

设置

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

# Setting random seed to obtain reproducible results.
import tensorflow as tf

tf.random.set_seed(42)

import keras
from keras import layers

import os
import glob
import imageio.v2 as imageio
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

# Initialize global variables.
AUTO = tf.data.AUTOTUNE
BATCH_SIZE = 5
NUM_SAMPLES = 32
POS_ENCODE_DIMS = 16
EPOCHS = 20

下载并加载数据

npz 数据文件包含图像、相机姿势和焦距。如下图所示，图像是从多个摄像机角度拍摄的。

（下图：多角度拍摄）

要在这种情况下理解摄像机的姿势，我们必须首先让自己认为，摄像机是真实世界和二维图像之间的映射。

请看下面的等式：

其中 x 是二维图像点，X 是三维世界点，P 是摄像机矩阵。P 是一个 3 x 4 矩阵，在将现实世界中的物体映射到图像平面上起着至关重要的作用。

摄像机矩阵是一个仿射变换矩阵，它与 3 x 1 列[图像高度、图像宽度、焦距]相连接，从而产生姿态矩阵。

该矩阵的尺寸为 3 x 5，其中第一个 3 x 3 块为摄像机视角。轴线为[向下、向右、向后]或[-y、x、z]，其中摄像机朝向前方-z。

（下图：仿射变换）

COLMAP 框架为 [右、下、前] 或 [x、-y、-z]。点击此处了解更多有关 COLMAP 的信息。

# Download the data if it does not already exist.
url = (
    "http://cseweb.ucsd.edu/~viscomp/projects/LF/papers/ECCV20/nerf/tiny_nerf_data.npz"
)
data = keras.utils.get_file(origin=url)

data = np.load(data)
images = data["images"]
im_shape = images.shape
(num_images, H, W, _) = images.shape
(poses, focal) = (data["poses"], data["focal"])

# Plot a random image from the dataset for visualization.
plt.imshow(images[np.random.randint(low=0, high=num_images)])
plt.show()

数据管道

在了解了摄像机矩阵的概念以及三维场景到二维图像的映射之后，我们来谈谈反向映射，即从二维图像到三维场景的映射。

我们需要讨论使用光线投射和追踪进行体积渲染的问题，这些都是常见的计算机图形技术。本节将帮助你快速掌握这些技术。

考虑一幅有 N 个像素的图像。我们用光线穿过每个像素，并对光线上的一些点进行采样。如下图所示，射线的参数方程通常为 r(t) = o + td，其中 t 为参数，o 为原点，d 为单位方向向量。

（下图：r(t) = o + td，其中 t 为 3）

在下图中，我们考虑了一条射线，并对射线上的一些随机点进行了采样。

这些采样点都有一个唯一的位置（x、y、z），射线有一个视角（theta、phi）。视角尤其有趣，因为我们可以用多种不同的方式将光线射穿一个像素，每种方式都有一个独特的视角。

另一个值得注意的地方是采样过程中添加的噪声。

我们在每个样本中都添加了均匀的噪声，这样样本就对应于一个连续的分布。在下图中，蓝色的点是均匀分布的样本，白色的点（t1、t2、t3）是随机放置在样本之间的。

（下图：从射线中采样点。）

下图展示了整个三维采样过程，可以看到射线从白色图像中射出。

这意味着每个像素都会有相应的光线，而每条光线都会在不同的点上采样。

（下图：从三维图像的所有像素射出光线）

这些采样点是 NeRF 模型的输入。然后要求模型预测该点的 RGB 颜色和体积密度。

（下图：数据管道）

def encode_position(x):
    """Encodes the position into its corresponding Fourier feature.

    Args:
        x: The input coordinate.

    Returns:
        Fourier features tensors of the position.
    """
    positions = [x]
    for i in range(POS_ENCODE_DIMS):
        for fn in [tf.sin, tf.cos]:
            positions.append(fn(2.0**i * x))
    return tf.concat(positions, axis=-1)


def get_rays(height, width, focal, pose):
    """Computes origin point and direction vector of rays.

    Args:
        height: Height of the image.
        width: Width of the image.
        focal: The focal length between the images and the camera.
        pose: The pose matrix of the camera.

    Returns:
        Tuple of origin point and direction vector for rays.
    """
    # Build a meshgrid for the rays.
    i, j = tf.meshgrid(
        tf.range(width, dtype=tf.float32),
        tf.range(height, dtype=tf.float32),
        indexing="xy",
    )

    # Normalize the x axis coordinates.
    transformed_i = (i - width * 0.5) / focal

    # Normalize the y axis coordinates.
    transformed_j = (j - height * 0.5) / focal

    # Create the direction unit vectors.
    directions = tf.stack([transformed_i, -transformed_j, -tf.ones_like(i)], axis=-1)

    # Get the camera matrix.
    camera_matrix = pose[:3, :3]
    height_width_focal = pose[:3, -1]

    # Get origins and directions for the rays.
    transformed_dirs = directions[..., None, :]
    camera_dirs = transformed_dirs * camera_matrix
    ray_directions = tf.reduce_sum(camera_dirs, axis=-1)
    ray_origins = tf.broadcast_to(height_width_focal, tf.shape(ray_directions))

    # Return the origins and directions.
    return (ray_origins, ray_directions)


def render_flat_rays(ray_origins, ray_directions, near, far, num_samples, rand=False):
    """Renders the rays and flattens it.

    Args:
        ray_origins: The origin points for rays.
        ray_directions: The direction unit vectors for the rays.
        near: The near bound of the volumetric scene.
        far: The far bound of the volumetric scene.
        num_samples: Number of sample points in a ray.
        rand: Choice for randomising the sampling strategy.

    Returns:
       Tuple of flattened rays and sample points on each rays.
    """
    # Compute 3D query points.
    # Equation: r(t) = o+td -> Building the "t" here.
    t_vals = tf.linspace(near, far, num_samples)
    if rand:
        # Inject uniform noise into sample space to make the sampling
        # continuous.
        shape = list(ray_origins.shape[:-1]) + [num_samples]
        noise = tf.random.uniform(shape=shape) * (far - near) / num_samples
        t_vals = t_vals + noise

    # Equation: r(t) = o + td -> Building the "r" here.
    rays = ray_origins[..., None, :] + (
        ray_directions[..., None, :] * t_vals[..., None]
    )
    rays_flat = tf.reshape(rays, [-1, 3])
    rays_flat = encode_position(rays_flat)
    return (rays_flat, t_vals)


def map_fn(pose):
    """Maps individual pose to flattened rays and sample points.

    Args:
        pose: The pose matrix of the camera.

    Returns:
        Tuple of flattened rays and sample points corresponding to the
        camera pose.
    """
    (ray_origins, ray_directions) = get_rays(height=H, width=W, focal=focal, pose=pose)
    (rays_flat, t_vals) = render_flat_rays(
        ray_origins=ray_origins,
        ray_directions=ray_directions,
        near=2.0,
        far=6.0,
        num_samples=NUM_SAMPLES,
        rand=True,
    )
    return (rays_flat, t_vals)


# Create the training split.
split_index = int(num_images * 0.8)

# Split the images into training and validation.
train_images = images[:split_index]
val_images = images[split_index:]

# Split the poses into training and validation.
train_poses = poses[:split_index]
val_poses = poses[split_index:]

# Make the training pipeline.
train_img_ds = tf.data.Dataset.from_tensor_slices(train_images)
train_pose_ds = tf.data.Dataset.from_tensor_slices(train_poses)
train_ray_ds = train_pose_ds.map(map_fn, num_parallel_calls=AUTO)
training_ds = tf.data.Dataset.zip((train_img_ds, train_ray_ds))
train_ds = (
    training_ds.shuffle(BATCH_SIZE)
    .batch(BATCH_SIZE, drop_remainder=True, num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

# Make the validation pipeline.
val_img_ds = tf.data.Dataset.from_tensor_slices(val_images)
val_pose_ds = tf.data.Dataset.from_tensor_slices(val_poses)
val_ray_ds = val_pose_ds.map(map_fn, num_parallel_calls=AUTO)
validation_ds = tf.data.Dataset.zip((val_img_ds, val_ray_ds))
val_ds = (
    validation_ds.shuffle(BATCH_SIZE)
    .batch(BATCH_SIZE, drop_remainder=True, num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

NeRF 模型

该模型为多层感知器（MLP），以 ReLU 作为非线性。

某论文节选：

"我们通过限制网络仅预测体积密度 sigma 与位置 x 的函数关系，同时允许预测 RGB 颜色 c 与位置和观察方向的函数关系，从而使表示法与多视角一致。为此，MLP 首先使用 8 个全连接层（每层使用 ReLU 激活和 256 个通道）处理输入的 3D 坐标 x，并输出 sigma 和 256 维特征向量。然后，将该特征向量与摄像机光线的观察方向连接起来，并传递给另外一个全连接层（使用 ReLU 激活和 128 个通道），该层输出与视图相关的 RGB 颜色"。

在这里，我们采用了最小化的实现方式，使用了 64 个密集单元，而不是论文中提到的 256 个。

def get_nerf_model(num_layers, num_pos):
    """Generates the NeRF neural network.

    Args:
        num_layers: The number of MLP layers.
        num_pos: The number of dimensions of positional encoding.

    Returns:
        The `keras` model.
    """
    inputs = keras.Input(shape=(num_pos, 2 * 3 * POS_ENCODE_DIMS + 3))
    x = inputs
    for i in range(num_layers):
        x = layers.Dense(units=64, activation="relu")(x)
        if i % 4 == 0 and i > 0:
            # Inject residual connection.
            x = layers.concatenate([x, inputs], axis=-1)
    outputs = layers.Dense(units=4)(x)
    return keras.Model(inputs=inputs, outputs=outputs)


def render_rgb_depth(model, rays_flat, t_vals, rand=True, train=True):
    """Generates the RGB image and depth map from model prediction.

    Args:
        model: The MLP model that is trained to predict the rgb and
            volume density of the volumetric scene.
        rays_flat: The flattened rays that serve as the input to
            the NeRF model.
        t_vals: The sample points for the rays.
        rand: Choice to randomise the sampling strategy.
        train: Whether the model is in the training or testing phase.

    Returns:
        Tuple of rgb image and depth map.
    """
    # Get the predictions from the nerf model and reshape it.
    if train:
        predictions = model(rays_flat)
    else:
        predictions = model.predict(rays_flat)
    predictions = tf.reshape(predictions, shape=(BATCH_SIZE, H, W, NUM_SAMPLES, 4))

    # Slice the predictions into rgb and sigma.
    rgb = tf.sigmoid(predictions[..., :-1])
    sigma_a = tf.nn.relu(predictions[..., -1])

    # Get the distance of adjacent intervals.
    delta = t_vals[..., 1:] - t_vals[..., :-1]
    # delta shape = (num_samples)
    if rand:
        delta = tf.concat(
            [delta, tf.broadcast_to([1e10], shape=(BATCH_SIZE, H, W, 1))], axis=-1
        )
        alpha = 1.0 - tf.exp(-sigma_a * delta)
    else:
        delta = tf.concat(
            [delta, tf.broadcast_to([1e10], shape=(BATCH_SIZE, 1))], axis=-1
        )
        alpha = 1.0 - tf.exp(-sigma_a * delta[:, None, None, :])

    # Get transmittance.
    exp_term = 1.0 - alpha
    epsilon = 1e-10
    transmittance = tf.math.cumprod(exp_term + epsilon, axis=-1, exclusive=True)
    weights = alpha * transmittance
    rgb = tf.reduce_sum(weights[..., None] * rgb, axis=-2)

    if rand:
        depth_map = tf.reduce_sum(weights * t_vals, axis=-1)
    else:
        depth_map = tf.reduce_sum(weights * t_vals[:, None, None], axis=-1)
    return (rgb, depth_map)

训练

训练步骤是作为自定义 keras.Model 子类的一部分实现的，这样我们就可以使用 model.fit 功能。

class NeRF(keras.Model):
    def __init__(self, nerf_model):
        super().__init__()
        self.nerf_model = nerf_model

    def compile(self, optimizer, loss_fn):
        super().compile()
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.psnr_metric = keras.metrics.Mean(name="psnr")

    def train_step(self, inputs):
        # Get the images and the rays.
        (images, rays) = inputs
        (rays_flat, t_vals) = rays

        with tf.GradientTape() as tape:
            # Get the predictions from the model.
            rgb, _ = render_rgb_depth(
                model=self.nerf_model, rays_flat=rays_flat, t_vals=t_vals, rand=True
            )
            loss = self.loss_fn(images, rgb)

        # Get the trainable variables.
        trainable_variables = self.nerf_model.trainable_variables

        # Get the gradeints of the trainiable variables with respect to the loss.
        gradients = tape.gradient(loss, trainable_variables)

        # Apply the grads and optimize the model.
        self.optimizer.apply_gradients(zip(gradients, trainable_variables))

        # Get the PSNR of the reconstructed images and the source images.
        psnr = tf.image.psnr(images, rgb, max_val=1.0)

        # Compute our own metrics
        self.loss_tracker.update_state(loss)
        self.psnr_metric.update_state(psnr)
        return {"loss": self.loss_tracker.result(), "psnr": self.psnr_metric.result()}

    def test_step(self, inputs):
        # Get the images and the rays.
        (images, rays) = inputs
        (rays_flat, t_vals) = rays

        # Get the predictions from the model.
        rgb, _ = render_rgb_depth(
            model=self.nerf_model, rays_flat=rays_flat, t_vals=t_vals, rand=True
        )
        loss = self.loss_fn(images, rgb)

        # Get the PSNR of the reconstructed images and the source images.
        psnr = tf.image.psnr(images, rgb, max_val=1.0)

        # Compute our own metrics
        self.loss_tracker.update_state(loss)
        self.psnr_metric.update_state(psnr)
        return {"loss": self.loss_tracker.result(), "psnr": self.psnr_metric.result()}

    @property
    def metrics(self):
        return [self.loss_tracker, self.psnr_metric]


test_imgs, test_rays = next(iter(train_ds))
test_rays_flat, test_t_vals = test_rays

loss_list = []


class TrainMonitor(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        loss = logs["loss"]
        loss_list.append(loss)
        test_recons_images, depth_maps = render_rgb_depth(
            model=self.model.nerf_model,
            rays_flat=test_rays_flat,
            t_vals=test_t_vals,
            rand=True,
            train=False,
        )

        # Plot the rgb, depth and the loss plot.
        fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(20, 5))
        ax[0].imshow(keras.utils.array_to_img(test_recons_images[0]))
        ax[0].set_title(f"Predicted Image: {epoch:03d}")

        ax[1].imshow(keras.utils.array_to_img(depth_maps[0, ..., None]))
        ax[1].set_title(f"Depth Map: {epoch:03d}")

        ax[2].plot(loss_list)
        ax[2].set_xticks(np.arange(0, EPOCHS + 1, 5.0))
        ax[2].set_title(f"Loss Plot: {epoch:03d}")

        fig.savefig(f"images/{epoch:03d}.png")
        plt.show()
        plt.close()


num_pos = H * W * NUM_SAMPLES
nerf_model = get_nerf_model(num_layers=8, num_pos=num_pos)

model = NeRF(nerf_model)
model.compile(
    optimizer=keras.optimizers.Adam(), loss_fn=keras.losses.MeanSquaredError()
)

# Create a directory to save the images during training.
if not os.path.exists("images"):
    os.makedirs("images")

model.fit(
    train_ds,
    validation_data=val_ds,
    batch_size=BATCH_SIZE,
    epochs=EPOCHS,
    callbacks=[TrainMonitor()],
)


def create_gif(path_to_images, name_gif):
    filenames = glob.glob(path_to_images)
    filenames = sorted(filenames)
    images = []
    for filename in tqdm(filenames):
        images.append(imageio.imread(filename))
    kargs = {"duration": 0.25}
    imageio.mimsave(name_gif, images, "GIF", **kargs)


create_gif("images/*.png", "training.gif")

演绎展示：

Epoch 1/20
  1/16 ━[37m━━━━━━━━━━━━━━━━━━━  3:54 16s/step - loss: 0.0948 - psnr: 10.6234

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1699908753.457905   65271 device_compiler.h:187] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.

 1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 924ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 29s 889ms/step - loss: 0.1091 - psnr: 9.8283 - val_loss: 0.0753 - val_psnr: 11.5686
Epoch 2/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 477ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 926ms/step - loss: 0.0633 - psnr: 12.4819 - val_loss: 0.0657 - val_psnr: 12.1781
Epoch 3/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 921ms/step - loss: 0.0589 - psnr: 12.6268 - val_loss: 0.0637 - val_psnr: 12.3413
Epoch 4/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 470ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 915ms/step - loss: 0.0573 - psnr: 12.8150 - val_loss: 0.0617 - val_psnr: 12.4789
Epoch 5/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 477ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 918ms/step - loss: 0.0552 - psnr: 12.9703 - val_loss: 0.0594 - val_psnr: 12.6457
Epoch 6/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 476ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 894ms/step - loss: 0.0538 - psnr: 13.0895 - val_loss: 0.0533 - val_psnr: 13.0049
Epoch 7/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 940ms/step - loss: 0.0436 - psnr: 13.9857 - val_loss: 0.0381 - val_psnr: 14.4764
Epoch 8/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 475ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 919ms/step - loss: 0.0325 - psnr: 15.1856 - val_loss: 0.0294 - val_psnr: 15.5187
Epoch 9/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 478ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 927ms/step - loss: 0.0276 - psnr: 15.8105 - val_loss: 0.0259 - val_psnr: 16.0297
Epoch 10/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 952ms/step - loss: 0.0251 - psnr: 16.1994 - val_loss: 0.0252 - val_psnr: 16.0842
Epoch 11/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 909ms/step - loss: 0.0239 - psnr: 16.3749 - val_loss: 0.0228 - val_psnr: 16.5269
Epoch 12/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 19s 1s/step - loss: 0.0215 - psnr: 16.8117 - val_loss: 0.0186 - val_psnr: 17.3930
Epoch 13/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 923ms/step - loss: 0.0188 - psnr: 17.3916 - val_loss: 0.0174 - val_psnr: 17.6570
Epoch 14/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 476ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 973ms/step - loss: 0.0175 - psnr: 17.6871 - val_loss: 0.0172 - val_psnr: 17.6644
Epoch 15/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 468ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 919ms/step - loss: 0.0172 - psnr: 17.7639 - val_loss: 0.0161 - val_psnr: 18.0313
Epoch 16/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 477ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 915ms/step - loss: 0.0150 - psnr: 18.3860 - val_loss: 0.0151 - val_psnr: 18.2832
Epoch 17/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 926ms/step - loss: 0.0154 - psnr: 18.2210 - val_loss: 0.0146 - val_psnr: 18.4284
Epoch 18/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 468ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 959ms/step - loss: 0.0145 - psnr: 18.4869 - val_loss: 0.0134 - val_psnr: 18.8039
Epoch 19/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 16s 933ms/step - loss: 0.0136 - psnr: 18.8040 - val_loss: 0.0138 - val_psnr: 18.6680
Epoch 20/20
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 472ms/step

 16/16 ━━━━━━━━━━━━━━━━━━━━ 15s 916ms/step - loss: 0.0131 - psnr: 18.9661 - val_loss: 0.0132 - val_psnr: 18.8687

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 59.40it/s]

可视化训练步骤

在这里，我们看到了训练步骤。随着损失的减少，渲染图像和深度图的效果越来越好。在本地系统中，你将看到生成的 training.gif 文件。

推理

在本节中，我们要求模型构建新的场景视图。在训练步骤中，我们给模型提供了 106 个场景视图。训练图像集合不可能包含场景的每个角度。经过训练的模型可以用一组稀疏的训练图像来表示整个三维场景。

在这里，我们向模型提供不同的姿势，并要求它给出与该摄像机视角相对应的二维图像。如果我们推断出所有 360 度视角的模型，那么它就能从四周提供整个场景的概览。

# Get the trained NeRF model and infer.
nerf_model = model.nerf_model
test_recons_images, depth_maps = render_rgb_depth(
    model=nerf_model,
    rays_flat=test_rays_flat,
    t_vals=test_t_vals,
    rand=True,
    train=False,
)

# Create subplots.
fig, axes = plt.subplots(nrows=5, ncols=3, figsize=(10, 20))

for ax, ori_img, recons_img, depth_map in zip(
    axes, test_imgs, test_recons_images, depth_maps
):
    ax[0].imshow(keras.utils.array_to_img(ori_img))
    ax[0].set_title("Original")

    ax[1].imshow(keras.utils.array_to_img(recons_img))
    ax[1].set_title("Reconstructed")

    ax[2].imshow(keras.utils.array_to_img(depth_map[..., None]), cmap="inferno")
    ax[2].set_title("Depth Map")

演绎展示：

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 475ms/step

渲染三维场景

在这里，我们将合成新颖的 3D 视图，并将所有视图拼接在一起，以渲染包含 360 度视图的视频。

def get_translation_t(t):
    """Get the translation matrix for movement in t."""
    matrix = [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, t],
        [0, 0, 0, 1],
    ]
    return tf.convert_to_tensor(matrix, dtype=tf.float32)


def get_rotation_phi(phi):
    """Get the rotation matrix for movement in phi."""
    matrix = [
        [1, 0, 0, 0],
        [0, tf.cos(phi), -tf.sin(phi), 0],
        [0, tf.sin(phi), tf.cos(phi), 0],
        [0, 0, 0, 1],
    ]
    return tf.convert_to_tensor(matrix, dtype=tf.float32)


def get_rotation_theta(theta):
    """Get the rotation matrix for movement in theta."""
    matrix = [
        [tf.cos(theta), 0, -tf.sin(theta), 0],
        [0, 1, 0, 0],
        [tf.sin(theta), 0, tf.cos(theta), 0],
        [0, 0, 0, 1],
    ]
    return tf.convert_to_tensor(matrix, dtype=tf.float32)


def pose_spherical(theta, phi, t):
    """
    Get the camera to world matrix for the corresponding theta, phi
    and t.
    """
    c2w = get_translation_t(t)
    c2w = get_rotation_phi(phi / 180.0 * np.pi) @ c2w
    c2w = get_rotation_theta(theta / 180.0 * np.pi) @ c2w
    c2w = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) @ c2w
    return c2w


rgb_frames = []
batch_flat = []
batch_t = []

# Iterate over different theta value and generate scenes.
for index, theta in tqdm(enumerate(np.linspace(0.0, 360.0, 120, endpoint=False))):
    # Get the camera to world matrix.
    c2w = pose_spherical(theta, -30.0, 4.0)

    #
    ray_oris, ray_dirs = get_rays(H, W, focal, c2w)
    rays_flat, t_vals = render_flat_rays(
        ray_oris, ray_dirs, near=2.0, far=6.0, num_samples=NUM_SAMPLES, rand=False
    )

    if index % BATCH_SIZE == 0 and index > 0:
        batched_flat = tf.stack(batch_flat, axis=0)
        batch_flat = [rays_flat]

        batched_t = tf.stack(batch_t, axis=0)
        batch_t = [t_vals]

        rgb, _ = render_rgb_depth(
            nerf_model, batched_flat, batched_t, rand=False, train=False
        )

        temp_rgb = [np.clip(255 * img, 0.0, 255.0).astype(np.uint8) for img in rgb]

        rgb_frames = rgb_frames + temp_rgb
    else:
        batch_flat.append(rays_flat)
        batch_t.append(t_vals)

rgb_video = "rgb_video.mp4"
imageio.mimwrite(rgb_video, rgb_frames, fps=30, quality=7, macro_block_size=None)

演绎展示：

1it [00:01,  1.02s/it]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 475ms/step

6it [00:03,  1.95it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 478ms/step

11it [00:05,  2.11it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

16it [00:07,  2.17it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 477ms/step

25it [00:10,  3.05it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 477ms/step

27it [00:12,  2.14it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 479ms/step

31it [00:14,  2.02it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 472ms/step

36it [00:16,  2.11it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

41it [00:18,  2.16it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 472ms/step

46it [00:21,  2.19it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 475ms/step

51it [00:23,  2.22it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

56it [00:25,  2.24it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 464ms/step

61it [00:27,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

66it [00:29,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 476ms/step

71it [00:32,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

76it [00:34,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 475ms/step

81it [00:36,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

86it [00:38,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 476ms/step

91it [00:40,  2.26it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 465ms/step

96it [00:43,  2.27it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

101it [00:45,  2.28it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

106it [00:47,  2.28it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 473ms/step

111it [00:49,  2.27it/s]

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 474ms/step

120it [00:52,  2.31it/s]
[swscaler @ 0x67626c0] Warning: data is not aligned! This can lead to a speed loss