Spaces:
Configuration error
Configuration error
File size: 1,694 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 |
import torch
from lib.config import cfg
from .nerf_net_utils import *
from .. import embedder
import numpy as np
import mcubes
import trimesh
from . import if_clight_renderer
class Renderer(if_clight_renderer.Renderer):
def __init__(self, net):
super(Renderer, self).__init__(net)
def batchify_rays(self, wpts, alpha_decoder, chunk=1024 * 32):
"""Render rays in smaller minibatches to avoid OOM.
"""
n_batch, n_point = wpts.shape[:2]
all_ret = []
for i in range(0, n_point, chunk):
ret = alpha_decoder(wpts[:, i:i + chunk])
all_ret.append(ret)
all_ret = torch.cat(all_ret, 1)
return all_ret
def render(self, batch):
pts = batch['pts']
sh = pts.shape
inside = batch['inside'][0].bool()
pts = pts[0][inside][None]
# encode neural body
sp_input = self.prepare_sp_input(batch)
feature_volume = self.net.encode_sparse_voxels(sp_input)
alpha_decoder = lambda x: self.net.calculate_density(
x, feature_volume, sp_input)
alpha = self.batchify_rays(pts, alpha_decoder, 2048 * 64)
alpha = alpha[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)
# vertices = (vertices - 10) * 0.005
# vertices = vertices + batch['wbounds'][0, 0].detach().cpu().numpy()
mesh = trimesh.Trimesh(vertices, triangles)
ret = {'cube': cube, 'mesh': mesh}
return ret
|