render.py

from functools import partial
import math
from jax import dlpack, lax, numpy as jnp
import jax
import toml

from src import jax_nerf
from src.common import NerfSynthetic
from src.model import NGPTrainer
from src.common import MIN_CONE_STEPSIZE, NERF_CASCADES

from scripts.common import *

import os
import numpy as np
import imageio.v3 as imageio

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

resolution = [800, 800]

# dataset_dir = "../baseline/data/nerf_synthetic"
dataset_dir = "../nerf_synthetic"
dataset_name = "chair"
dataset = NerfSynthetic(f"{dataset_dir}/{dataset_name}")

x = ((jnp.linspace(0, resolution[0], resolution[0], endpoint=False) + 0.5) /
     resolution[0] - dataset.principal[0]) * resolution[0] / dataset.focal[0]
y = ((jnp.linspace(0, resolution[1], resolution[1], endpoint=False) + 0.5) /
     resolution[1] - dataset.principal[1]) * resolution[1] / dataset.focal[1]

xv, yv = jnp.meshgrid(x, y)
xv = jnp.ravel(xv)
yv = jnp.ravel(yv)


@partial(jax.vmap, in_axes=(0, 0, None))
def get_ray(xv, yv, transform_matrix):
    ray_o = transform_matrix[:, 3]
    ray_d = jnp.array([xv, yv, 1.0])
    ray_d = jnp.matmul(transform_matrix[:3, :3], ray_d)
    ray_d = ray_d / jnp.linalg.norm(ray_d)

    return jnp.array([ray_o, ray_d])


ray = get_ray(xv, yv, dataset.transform_matrix[0])
alive = jnp.ones((resolution[0] * resolution[1], ), dtype=jnp.bool_)
t = jnp.zeros((resolution[0] * resolution[1], ))
coords = jnp.zeros((resolution[0] * resolution[1], 7))

snapshot = NGPTrainer.load_msgpack(f"baseline/{dataset_name}.msgpack")
config = toml.load("configs/ngp/base.toml")
trainer = NGPTrainer(
    {
        "config": config,
        "density_grid": snapshot["density_grid"],
        "params": snapshot["params"]
    }, dataset)

ray = dlpack.to_dlpack(ray)
density_grid = snapshot["density_grid"]
t = dlpack.to_dlpack(t)


def unwarp_dt(dt):
    max_stepsize = MIN_CONE_STEPSIZE() * (1 << (NERF_CASCADES() - 1))
    return dt * (max_stepsize - MIN_CONE_STEPSIZE()) + MIN_CONE_STEPSIZE()


def logistic(x):
    return 1.0 / (1.0 + jnp.exp(-x))


def composite(local_network_output, local_rgba, local_dt):
    T = 1. - local_rgba[3]
    dt = unwarp_dt(local_dt)
    alpha = 1. - jnp.exp(-jnp.exp(local_network_output[3]) * dt)
    rgb = logistic(local_network_output[0:3])
    weight = alpha * T
    return jnp.append(rgb * weight, weight)


def composite_or_not(local_network_output, alive, local_rgba, local_dt):
    return lax.cond(alive, composite, lambda x, y, z: jnp.zeros(shape=(4, )),
                    local_network_output, local_rgba, local_dt)


rgba = jnp.zeros(shape=(resolution[0] * resolution[1], 4), dtype=jnp.float32)


def early_termination(local_rgba):
    return local_rgba[3] <= 1 - 1e-4


while jnp.count_nonzero(alive) != 0:

    coords = dlpack.to_dlpack(coords, True)
    alive = dlpack.to_dlpack(alive, True)
    jax_nerf.ngp_advance_pos_nerf(ray, density_grid, alive, t, coords,
                                  [dataset.aabb.min, dataset.aabb.max],
                                  resolution[0] * resolution[1])

    alive = dlpack.from_dlpack(alive)
    coords = dlpack.from_dlpack(coords)
    network_output = trainer.model(trainer.params, coords)
    partial_result = jax.vmap(composite_or_not,
                              in_axes=(0, 0, 0, 0))(network_output, alive,
                                                    rgba,
                                                    jnp.swapaxes(coords, 0,
                                                                 1)[3])
    rgba += partial_result
    ea = jax.vmap(early_termination)(rgba)
    alive = jnp.logical_and(ea, alive)
    print(jnp.count_nonzero(alive))

arrOut = np.array(rgba)
arrOut = np.uint8(arrOut * 255.0)
arrOut = np.reshape(arrOut, (resolution[1], resolution[0], 4))
imageio.imwrite("a.png", arrOut)

ref = np.array(dataset.images[0])
ref_image = ref.astype(np.float32) / 255.0

ref_image[..., :3] = srgb_to_linear(ref_image[..., :3])
ref_image[..., :3] *= ref_image[..., 3:4]
ref_image[..., :3] = linear_to_srgb(ref_image[..., :3])

rgba = np.reshape(np.array(rgba), (resolution[0], resolution[1], 4))

diff = np.abs(rgba[..., :3] - ref_image[..., :3]) * 255.0
diff = np.uint8(diff)
imageio.imwrite("diff.png", diff)

rgba = np.clip(rgba[..., :3], 0.0, 1.0)
ref_image = np.clip(ref_image[..., :3], 0.0, 1.0)
mse = float(compute_error("MSE", rgba, ref_image))
psnr = mse2psnr(mse)
print(psnr)