Logotipo de Zephyrnet

Entrenamiento de modelos Neural Radiance Field (NeRF) con Keras/TensorFlow y DeepVision

Fecha:

Campos de radiación neuronal, coloquialmente conocidos como NeRF han golpeado al mundo por sorpresa en 2020, publicado junto con el papel "NeRF: representación de escenas como campos de radiación neuronal para la síntesis de vistas", y siguen siendo la piedra angular de la síntesis de alta calidad de vistas novedosas, con imágenes y posiciones de cámara escasas.

Desde entonces, han encontrado numerosas aplicaciones, pero probablemente más destacadas en el modelado volumétrico geoespacial, con empresas como Google que confían en NeRF para crear estructuras 3D de edificios y sitios patrimoniales desde varios ángulos de imágenes satelitales, y empresas especializadas en realizar reconstrucción 3D y digitalización de sitios culturales bien conocidos.

En esta guía, entrenaremos un modelo Neural Radiance Field (NeRF) en el conjunto de datos Tiny NeRF original, usando TensorFlow/Keras y DeepVision, para realizar una nueva síntesis de vista/reconstrucción 3D.

En una sola hora, en una máquina comercial, generará vistas novedosas de imágenes del conjunto de datos TinyNeRF:

Síntesis de vista novedosa y campos de radiación neuronal

Esta sección proporciona una introducción/resumen simplificado sobre la forma en que funcionan los campos de radiación neuronal, pero puede llevar algún tiempo comprender de forma intuitiva cómo funcionan si es nuevo en el campo.

Nota: El documento original, así como el video educativo y los gráficos asociados con él, son excelentes materiales de aprendizaje. Si está interesado en comprender el concepto subyacente de campos de radiación en los que se basan los NeRF para representar una escena, el Entrada de Wikipedia para "campos de luz" proporciona una gran introducción, pero se pueden resumir de una manera de alto nivel como

“El campo de luz es una función vectorial que describe la cantidad de luz que fluye en cada dirección a través de cada punto del espacio”.

Los NeRF se utilizan para síntesis de vista de novela – crear nuevas vistas de objetos e imágenes, dadas algunas vistas. En efecto, puede pensar en la síntesis de vista novedosa como una conversión 2D->3D, y existen muchos enfoques para resolver este problema, algunos más exitosos que otros.

Históricamente un problema desafiante, la solución propuesta por NeRFs es extremadamente simple pero produce resultados de vanguardia, generando imágenes de muy alta calidad desde ángulos novedosos:

Esto, naturalmente, los posicionó como un enfoque fundamental para resolver la síntesis de vistas novedosas, con muchos artículos posteriores que exploran, ajustan y mejoran las ideas presentes en ellos.

Consejo: El página web publicado junto con el documento contiene una muestra asombrosa del método y sus resultados, y un video educativo que construye una buena intuición de cómo funcionan estas redes ha sido lanzado oficialmente.

La canalización de los datos a los resultados se puede resumir como:

Donde la red neuronal aprende de imágenes dispersas con generación sintética rayos que se proyectan y muestrean a intervalos regulares. Las imágenes se colocan en el espacio dados los metadatos sobre las imágenes, como las posiciones de la cámara cuando se tomaron las imágenes. Debido a esto, no puede simplemente ingresar cualquier imagen y exigir posiciones de la cámara para poder posicionar con precisión las imágenes en el espacio para que los rayos creen un conjunto comprensible de puntos. Los puntos muestreados luego forman un conjunto 3D de puntos que representan la escena volumétrica:

La red neuronal se aproxima a un función de escena volumétrica – los valores RGB y la densidad (σ) de una escena. En efecto, entrenamos la red para memorizar el color y la densidad de cada punto de entrada, para poder reconstruir las imágenes desde puntos de vista novedosos. Dicho esto, los NeRF no están entrenados en un conjunto de imágenes y pueden extrapolarse a otras nuevas. Los NeRF están entrenados para codificar una escena, y luego solo se usan para esa escena, como los pesos de la propia red representan la escena.

Este es el principal "inconveniente" de los NeRF: debe entrenar una red para cada escena que desee codificar, y el proceso de entrenamiento es algo lento y requiere mucha memoria para entradas grandes. Las mejoras en el tiempo de entrenamiento son un área de investigación, con técnicas novedosas como “Optimización directa de cuadrícula de vóxeles” que mejoran significativamente el tiempo de entrenamiento sin sacrificar la calidad de imagen en el proceso.

Campos de radiación neuronal en DeepVision y TensorFlow

Las implementaciones de NeRF pueden ser un poco abrumadoras para aquellos que son nuevos en la representación volumétrica, y los repositorios de código suelen incluir muchos métodos auxiliares para manejar datos volumétricos, lo que puede parecer poco intuitivo para algunos. DeepVision es una biblioteca de visión por computadora novedosa que tiene como objetivo unificar la visión por computadora bajo una API común, con backends intercambiables (TensorFlow y PyTorch), conversiones de peso automáticas entre modelos y modelos con implementaciones idénticas en marcos de backend.

Para bajar la barrera de entrada, visión profunda ofrece una implementación simple pero fiel al original de los modelos Neural Radiance Field, con múltiples configuraciones para adaptarse a máquinas más y menos potentes con diferentes configuraciones de hardware:

  • NeRFTiny
  • NeRFSmall
  • NeRFMedium
  • NeRF
  • NeRFLarge

Se utilizan dos parámetros para crear estas configuraciones: width y depth. Dado que los NeRF son, en esencia, solo un modelo MLP que consta de tf.keras.layers.Dense() capas (con una sola concatenación entre capas), el depth representa directamente el número de Dense capas, mientras width representa el número de unidades utilizadas en cada uno.

NeRF corresponde a la configuración utilizada en el documento original, pero puede ser difícil de ejecutar en algunas máquinas locales, en cuyo caso, NeRFMedium proporciona un rendimiento muy similar con requisitos de memoria más pequeños.

Prosigamos e instalemos DeepVision con pip:

$ pip install deepvision-toolkit

Instanciar un modelo es tan fácil como:

import deepvision
model = deepvision.models.NeRFMedium(input_shape=(num_pos, input_features), backend='tensorflow') model.summary()

El modelo en sí es extremadamente simple:

Model: "ne_rftf"
__________________________________________________________________________________________________ Layer (type) Output Shape Param # Connected to ================================================================================================== input_1 (InputLayer) [(None, 640000, 195 0 [] )] dense (Dense) (None, 640000, 128) 25088 ['input_1[0][0]'] dense_1 (Dense) (None, 640000, 128) 16512 ['dense[0][0]'] dense_2 (Dense) (None, 640000, 128) 16512 ['dense_1[0][0]'] dense_3 (Dense) (None, 640000, 128) 16512 ['dense_2[0][0]'] dense_4 (Dense) (None, 640000, 128) 16512 ['dense_3[0][0]'] concatenate (Concatenate) (None, 640000, 323) 0 ['dense_4[0][0]', 'input_1[0][0]'] dense_5 (Dense) (None, 640000, 128) 41472 ['concatenate[0][0]'] dense_6 (Dense) (None, 640000, 4) 516 ['dense_5[0][0]'] ==================================================================================================
Total params: 133,128
Trainable params: 133,124
Non-trainable params: 4
__________________________________________________________________________________________________

Consulte nuestra guía práctica y práctica para aprender Git, con las mejores prácticas, los estándares aceptados por la industria y la hoja de trucos incluida. Deja de buscar en Google los comandos de Git y, de hecho, aprenden ella!

Veremos más de cerca cómo manejar los resultados del modelo y cómo renderizar las imágenes producidas por los pesos del modelo, en un momento.

Cargando el conjunto de datos TinyNeRF

Dado que los NeRF pueden ser algo costosos de entrenar en imágenes de entrada más grandes, se lanzaron con un pequeño conjunto de datos de 100 × 100 imágenes, denominado TinyNeRF para facilitar las pruebas y las iteraciones. Posteriormente, se ha convertido en un conjunto de datos clásico para probar NeRF y para ingresar al campo, de manera similar a cómo MNIST se convirtió en el "Hola mundo" del reconocimiento de dígitos.

El conjunto de datos está disponible como .npz archivo, y contiene imágenes, puntos focales (utilizados para la normalización) y poses de cámara, y se puede obtener del lanzamiento del código oficial:

import requests
import numpy as np
import matplotlib.pyplot as plt url = "https://people.eecs.berkeley.edu/~bmild/nerf/tiny_nerf_data.npz"
save_path = 'tiny_nerf.npz' file_data = requests.get(url).content
with open(save_path, "wb") as file: file.write(file_data) data = np.load(save_path) images, poses, focal = data["images"], data["poses"], data["focal"] print(images.shape) print(poses.shape) print(focal) 

Hay 106 imágenes, 100×100 cada una, con 3 canales (RGB). Todas las imágenes son de una pequeña excavadora de lego. Grafiquemos las primeras cinco imágenes:

fig, ax = plt.subplots(1, 5, figsize=(20, 12))
for i in range(5): ax[i].imshow(images[i])

El posiciones de la cámara proporcionados en el conjunto de datos son cruciales para poder reconstruir el espacio en el que se tomaron las imágenes, lo que nos permite proyectar rayos a través de las imágenes y formar un espacio volumétrico con los puntos muestreados en cada proyección.

Sin embargo, dado que este conjunto de datos requiere mucha preparación para la fase de entrenamiento, DeepVision ofrece una load_tiny_nerf() cargador de conjuntos de datos, que realizará la preparación por usted, con un opcional validation_split, pos_embed y num_ray_samples, y devuelve una vainilla tf.data.Dataset que puede crear canalizaciones de alto rendimiento con:

import deepvision train_ds, valid_ds = deepvision.datasets.load_tiny_nerf(pos_embed=16, num_ray_samples=32, save_path='tiny_nerf.npz', validation_split=0.2, backend='tensorflow') 

Absolutamente no necesita crear un conjunto de validación aquí, ya que el punto is para sobreajustar y memorizar completamente las imágenes, y el conjunto de validación aquí se crea principalmente como una verificación de cordura.

Echemos un vistazo a la longitud y las formas de entrada en el conjunto de datos de entrenamiento:

print('Train dataset length:', len(train_ds))
print(train_ds)

Esto resulta en:

Train dataset length: 84
<ZipDataset element_spec=(TensorSpec(shape=(100, 100, 3), dtype=tf.float32, name=None), (TensorSpec(shape=(320000, 99), dtype=tf.float32, name=None), TensorSpec(shape=(100, 100, 32), dtype=tf.float32, name=None)))>

El pos_embed argumento establece el número de incrustaciones posicionales se utiliza para transformar las coordenadas 5D (x, y, z y ángulos de visión Theta y Phi). Las incrustaciones posicionales fueron crucial para que la red pudiera representar funciones de mayor frecuencia, que era un "ingrediente faltante" para hacer que este tipo de técnica funcionara en el pasado, ya que las redes luchaban por aproximar funciones que representaban variaciones de alta frecuencia en color y geometría, debido a su sesgo hacia el aprendizaje de funciones de baja frecuencia en su lugar:

El num_ray_samples representa el número de muestras tomadas a lo largo de cada rayo proyectado en la imagen.

Naturalmente, cuantas más incrustaciones posicionales y muestras de rayos utilice, mayor será la resolución de la escena volumétrica que está aproximando y, por lo tanto, más detalladas serán las imágenes finales, a costa de mayores costos computacionales.

Entrenamiento de un NeRF con TensorFlow/Keras y DeepVision

Echemos un vistazo a un ejemplo completo de carga de datos, preparación del conjunto de datos, instanciación de un modelo y entrenamiento con DeepVision y el ecosistema TensorFlow/Keras:

import deepvision
from deepvision.datasets import load_tiny_nerf
import tensorflow as tf config = { 'img_height': 100, 'img_width': 100, 'pos_embed': 32, 'num_ray_samples': 64, 'batch_size': 1
} num_pos = config['img_height'] * config['img_width'] * config['num_ray_samples']
input_features = 6 * config['pos_embed'] + 3 train_ds, valid_ds = load_tiny_nerf(pos_embed=config['pos_embed'], num_ray_samples=config['num_ray_samples'], save_path='tiny_nerf.npz', validation_split=0.2, backend='tensorflow') train_ds = train_ds.batch(config['batch_size']).prefetch(tf.data.AUTOTUNE)
valid_ds = valid_ds.batch(config['batch_size']).prefetch(tf.data.AUTOTUNE) model = deepvision.models.NeRFMedium(input_shape=(num_pos, input_features), backend='tensorflow') model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), loss=tf.keras.losses.MeanSquaredError()) callbacks = [tf.keras.callbacks.ReduceLROnPlateau()] history = model.fit(train_ds, epochs=50, validation_data=valid_ds, callbacks=callbacks)

En una Nvidia GTX1660Super, el entrenamiento con 32 incrustaciones posicionales y 64 muestras de rayos toma ~1 minuto por época, pero las configuraciones más pequeñas, como 8-16 incrustaciones posicionales y 32 muestras de rayos, pueden tomar tan solo ~7 segundos por época:

Epoch 1/50
84/84 [==============================] - 65s 746ms/step - loss: 0.0603 - psnr: 12.6432 - val_loss: 0.0455 - val_psnr: 13.7601 - lr: 0.0010
...
Epoch 50/50
84/84 [==============================] - 55s 658ms/step - loss: 0.0039 - psnr: 24.1984 - val_loss: 0.0043 - val_psnr: 23.8576 - lr: 0.0010

Después de aproximadamente una hora, en una sola GPU comercial, el modelo alcanza ~24 PSNR. Lo que pasa con los NeRF es que cuanto más entrenes, más se acercará a las representaciones de las imágenes originales, lo que significa que, por lo general, verás que las métricas aumentan con el tiempo a medida que entrenas más. Ayuda tener un ReduceLROnPlateau devolución de llamada para manejar la reducción de la tasa de aprendizaje para ajustar los resultados cuando se acerca el final del entrenamiento.

El modelo informa dos métricas: loss y psnr. La pérdida es el error cuadrático medio de cada píxel y funciona como una gran función de pérdida para los NeRF, pero es difícil de interpretar.

Relación pico de señal a ruido (PSNR) es la relación entre la señal (potencia máxima de una señal) y el ruido (potencia del ruido que corrompe la fidelidad de la señal) que degrada la imagen. La relación pico de señal a ruido se puede utilizar como calidad de imagen métrico, y es muy intuitivo de interpretar para los humanos.

Ya con un PSNR de 24, las imágenes se vuelven bastante claras y los NeRF pueden alcanzar PSNR de más de 40 en TinyNeRF con suficiente tiempo de entrenamiento.

Visualización de salidas

La red genera un tensor de forma [batch_size, 640000, 4] donde los canales representan RGB y densidad, y los 640000 puntos codifican la escena. Para representarlos como imágenes, queremos remodelar el tensor a una forma de (batch_size, img_height, img_width, num_ray_samples, 4), y luego divida los 4 canales en RGB y sigma y procéselos en una imagen (y opcionalmente, un mapa de profundidad/precisión).

Específicamente, los canales RGB se pasan a través de un sigmoideo activación, mientras que el canal sigma pasa a través de un Rehacer activación, antes de ser procesada más y reducida a un tensor de forma (batch_size, img_height, img_width, rgb_channels), y dos tensores de forma (batch_size, img_height, img_width, depth_channel) y (batch_size, img_height, img_width, accuracy).

Para facilitar este proceso, podemos utilizar el nerf_render_image_and_depth_tf() funcionar desde volumetric_utils, que acepta el modelo para predecir RGB y sigma a partir de las entradas, y devuelve un lote de imágenes, mapas de profundidad y mapas de precisión:

import matplotlib.pyplot as plt
from deepvision.models.volumetric.volumetric_utils import nerf_render_image_and_depth_tf for batch in train_ds.take(5): (images, rays) = batch (rays_flat, t_vals) = rays image_batch, depth_maps, _ = nerf_render_image_and_depth_tf(model=model, rays_flat=rays_flat, t_vals=t_vals, img_height=config['img_height'], img_width=config['img_width'], num_ray_samples=config['num_ray_samples']) fig, ax = plt.subplots(1, 2) ax[0].imshow(tf.squeeze(image_batch[0])) ax[1].imshow(tf.squeeze(depth_maps[0]))

Aquí, estamos trazando 5 lotes (cada uno con una imagen) y sus mapas de profundidad.
Durante el entrenamiento, el modelo en sí se basa en el nerf_render_image_and_depth_tf() función para convertir predicciones en imágenes y calcular el error cuadrático medio y PSNR para los resultados. Ejecutar este código da como resultado:


Conclusión

En esta guía, hemos resumido algunos de los elementos clave de Neural Radiance Fields, como una breve introducción al tema, seguido de la carga y preparación del conjunto de datos TinyNeRF en TensorFlow, usando tf.datay entrenar un modelo NeRF con los ecosistemas Keras y DeepVision.

punto_img

Información más reciente

punto_img