Spaces:
Configuration error
Configuration error
File size: 3,942 Bytes
1ba539f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import torch
from lib.config import cfg
from .nerf_net_utils import *
import numpy as np
import mcubes
import trimesh
class Renderer:
def __init__(self, net):
self.net = net
def render_rays(self, ray_batch, net_c=None, pytest=False):
"""Volumetric rendering.
Args:
ray_batch: array of shape [batch_size, ...]. All information necessary
for sampling along a ray, including: ray origin, ray direction, min
dist, max dist, and unit-magnitude viewing direction.
network_fn: function. Model for predicting RGB and density at each point
in space.
network_query_fn: function used for passing queries to network_fn.
N_samples: int. Number of different times to sample along each ray.
retraw: bool. If True, include model's raw, unprocessed predictions.
lindisp: bool. If True, sample linearly in inverse depth rather than in depth.
perturb: float, 0 or 1. If non-zero, each ray is sampled at stratified
random points in time.
N_importance: int. Number of additional times to sample along each ray.
These samples are only passed to network_fine.
network_fine: "fine" network with same spec as network_fn.
white_bkgd: bool. If True, assume a white background.
raw_noise_std: ...
verbose: bool. If True, print more debugging info.
Returns:
rgb_map: [num_rays, 3]. Estimated RGB color of a ray. Comes from fine model.
disp_map: [num_rays]. Disparity map. 1 / depth.
acc_map: [num_rays]. Accumulated opacity along each ray. Comes from fine model.
raw: [num_rays, num_samples, 4]. Raw predictions from model.
rgb0: See rgb_map. Output for coarse model.
disp0: See disp_map. Output for coarse model.
acc0: See acc_map. Output for coarse model.
z_std: [num_rays]. Standard deviation of distances along ray for each
sample.
"""
pts = ray_batch
if net_c is None:
alpha = self.net(pts)
else:
alpha = self.net(pts, net_c)
if cfg.N_importance > 0:
alpha_0 = alpha
if net_c is None:
alpha = self.net(pts, model='fine')
else:
alpha = self.net(pts, net_c, model='fine')
ret = {
'alpha': alpha
}
if cfg.N_importance > 0:
ret['alpha0'] = alpha_0
for k in ret:
DEBUG = False
if (torch.isnan(ret[k]).any()
or torch.isinf(ret[k]).any()) and DEBUG:
print(f"! [Numerical Error] {k} contains nan or inf.")
return ret
def batchify_rays(self, rays_flat, chunk=1024 * 32):
"""Render rays in smaller minibatches to avoid OOM.
"""
all_ret = {}
for i in range(0, rays_flat.shape[0], chunk):
ret = self.render_rays(rays_flat[i:i + chunk])
for k in ret:
if k not in all_ret:
all_ret[k] = []
all_ret[k].append(ret[k])
all_ret = {k: torch.cat(all_ret[k], 0) for k in all_ret}
return all_ret
def render(self, batch):
pts = batch['pts']
sh = pts.shape
inside = batch['inside'][0].bool()
pts = pts[0][inside][None]
pts = pts.view(sh[0], -1, 1, 3)
ret = self.batchify_rays(pts, cfg.chunk)
alpha = ret['alpha']
alpha = alpha[0, :, 0, 0].detach().cpu().numpy()
cube = np.zeros(sh[1:-1])
inside = inside.detach().cpu().numpy()
cube[inside == 1] = alpha
cube = np.pad(cube, 10, mode='constant')
vertices, triangles = mcubes.marching_cubes(cube, cfg.mesh_th)
mesh = trimesh.Trimesh(vertices, triangles)
ret = {'cube': cube, 'mesh': mesh}
return ret
|