☕ reformat voxelisation
This commit is contained in:
parent
33adb79bf9
commit
67c023daba
1 changed files with 53 additions and 34 deletions
|
|
@ -7,8 +7,13 @@ import torch.distributed
|
|||
import trimesh
|
||||
from tqdm import tqdm
|
||||
|
||||
from ddp_train_nerf import (cleanup, config_parser, create_nerf, setup,
|
||||
setup_logger)
|
||||
from ddp_train_nerf import (
|
||||
cleanup,
|
||||
config_parser,
|
||||
create_nerf,
|
||||
setup,
|
||||
setup_logger,
|
||||
)
|
||||
|
||||
parser = config_parser()
|
||||
args = parser.parse_args()
|
||||
|
|
@ -21,13 +26,13 @@ args.rank = 0
|
|||
setup(args.rank, args.world_size)
|
||||
start, models = create_nerf(args.rank, args)
|
||||
|
||||
net_0 = models['net_0']
|
||||
net_0 = models["net_0"]
|
||||
|
||||
fg_far_depth = 1
|
||||
|
||||
# weird way to do it, should be change if something better exists
|
||||
for idx, m in enumerate(net_0.modules()):
|
||||
#print(idx, "->", m)
|
||||
# print(idx, "->", m)
|
||||
|
||||
# foreground
|
||||
if idx == 3:
|
||||
|
|
@ -48,7 +53,10 @@ for idx, m in enumerate(net_0.modules()):
|
|||
# put everything on GPU
|
||||
device = "cuda"
|
||||
|
||||
def query_occupancy(position, embedder_position, embedder_viewdir, mlp_net, device="cuda"):
|
||||
|
||||
def query_occupancy(
|
||||
position, embedder_position, embedder_viewdir, mlp_net, device="cuda"
|
||||
):
|
||||
"""
|
||||
Given a position returns the occupancy probabily of the network.
|
||||
|
||||
|
|
@ -80,7 +88,9 @@ def query_occupancy(position, embedder_position, embedder_viewdir, mlp_net, devi
|
|||
ray_d = ray_d / ray_d_norm
|
||||
|
||||
# forge the input
|
||||
nn_input = torch.cat((fg_embedder_position(position), fg_embedder_viewdir(ray_d)), dim=-1)
|
||||
nn_input = torch.cat(
|
||||
(fg_embedder_position(position), fg_embedder_viewdir(ray_d)), dim=-1
|
||||
)
|
||||
|
||||
# forward the NN
|
||||
nn_raw = mlp_net(nn_input)
|
||||
|
|
@ -88,18 +98,25 @@ def query_occupancy(position, embedder_position, embedder_viewdir, mlp_net, devi
|
|||
|
||||
return sigma
|
||||
|
||||
|
||||
# annonymous function
|
||||
f = lambda x, y, z: query_occupancy(torch.tensor([x,y,z], dtype=torch.float32, device=device), fg_embedder_position, fg_embedder_viewdir, mlp_net)
|
||||
f = lambda x, y, z: query_occupancy(
|
||||
torch.tensor([x, y, z], dtype=torch.float32, device=device),
|
||||
fg_embedder_position,
|
||||
fg_embedder_viewdir,
|
||||
mlp_net,
|
||||
)
|
||||
|
||||
|
||||
def marching_cube_and_render(sigma_list, threshold):
|
||||
|
||||
vertices, triangles = mcubes.marching_cubes(sigma_list, threshold)
|
||||
mesh = trimesh.Trimesh(vertices / N - .5, triangles)
|
||||
mesh = trimesh.Trimesh(vertices / N - 0.5, triangles)
|
||||
mesh.show()
|
||||
|
||||
|
||||
#position = torch.rand(3, device=device)
|
||||
#position = torch.tensor([0.1, 0.1, 0.1], device=device)
|
||||
# position = torch.rand(3, device=device)
|
||||
# position = torch.tensor([0.1, 0.1, 0.1], device=device)
|
||||
|
||||
ray_d = torch.rand(3, device=device)
|
||||
# normalize ray direction
|
||||
|
|
@ -108,47 +125,49 @@ ray_d = ray_d / ray_d_norm
|
|||
|
||||
|
||||
N = 100
|
||||
t = np.linspace(-1, 1, N+1)
|
||||
t = np.linspace(-1, 1, N + 1)
|
||||
|
||||
query_pts = np.stack(np.meshgrid(t, t, t), -1).astype(np.float32)
|
||||
#print(query_pts.shape)
|
||||
# print(query_pts.shape)
|
||||
sh = query_pts.shape
|
||||
flat = query_pts.reshape([-1,3])
|
||||
flat = query_pts.reshape([-1, 3])
|
||||
|
||||
#raw_voxel = torch.zeros(N+1, N+1, N+1, 4) # N, D, H, W
|
||||
fg_raw_voxel = torch.zeros(N+1, N+1, N+1)
|
||||
#bg_raw_voxel = torch.zeros(N+1, N+1, N+1)
|
||||
# raw_voxel = torch.zeros(N+1, N+1, N+1, 4) # N, D, H, W
|
||||
fg_raw_voxel = torch.zeros(N + 1, N + 1, N + 1)
|
||||
# bg_raw_voxel = torch.zeros(N+1, N+1, N+1)
|
||||
|
||||
i = 0
|
||||
for x,y,z in tqdm(flat):
|
||||
for x, y, z in tqdm(flat):
|
||||
|
||||
position = torch.tensor([x,y,z], device=device)
|
||||
#bg_position = torch.cat((position, torch.tensor([1], device=device)))
|
||||
position = torch.tensor([x, y, z], device=device)
|
||||
# bg_position = torch.cat((position, torch.tensor([1], device=device)))
|
||||
|
||||
# concat the output of the embedding
|
||||
fg_input = torch.cat((fg_embedder_position(position), fg_embedder_viewdir(ray_d)), dim=-1)
|
||||
#bg_input = torch.cat((bg_embedder_position(bg_position), bg_embedder_viewdir(ray_d)), dim=-1)
|
||||
fg_input = torch.cat(
|
||||
(fg_embedder_position(position), fg_embedder_viewdir(ray_d)), dim=-1
|
||||
)
|
||||
# bg_input = torch.cat((bg_embedder_position(bg_position), bg_embedder_viewdir(ray_d)), dim=-1)
|
||||
|
||||
# forward
|
||||
fg_raw = fg_mlp_net(fg_input)
|
||||
#bg_raw = bg_mlp_net(bg_input)
|
||||
# bg_raw = bg_mlp_net(bg_input)
|
||||
|
||||
#raw_voxel.append(position + float(nn_raw['sigma']))
|
||||
# raw_voxel.append(position + float(nn_raw['sigma']))
|
||||
fg_sigma = float(fg_raw["sigma"])
|
||||
#bg_sigma = float(bg_raw["sigma"])
|
||||
# bg_sigma = float(bg_raw["sigma"])
|
||||
|
||||
nx, ny, nz = np.unravel_index(i, (N+1, N+1, N+1))
|
||||
nx, ny, nz = np.unravel_index(i, (N + 1, N + 1, N + 1))
|
||||
i += 1 # update index
|
||||
#raw_voxel[unraveled_index] = torch.tensor([sigma, x, y, z])
|
||||
# raw_voxel[unraveled_index] = torch.tensor([sigma, x, y, z])
|
||||
fg_raw_voxel[nx, ny, nz] = fg_sigma
|
||||
#bg_raw_voxel[nx, ny, nz] = bg_sigma
|
||||
# bg_raw_voxel[nx, ny, nz] = bg_sigma
|
||||
|
||||
|
||||
fg_sigma = np.array(fg_raw_voxel)
|
||||
#bg_sigma = np.array(bg_raw_voxel)
|
||||
# bg_sigma = np.array(bg_raw_voxel)
|
||||
threshold = 0.5
|
||||
|
||||
|
||||
#vertices, triangles = mcubes.marching_cubes(sigma, threshold)
|
||||
#mesh = trimesh.Trimesh(vertices / N - .5, triangles)
|
||||
#mesh.show()
|
||||
# vertices, triangles = mcubes.marching_cubes(sigma, threshold)
|
||||
# mesh = trimesh.Trimesh(vertices / N - .5, triangles)
|
||||
# mesh.show()
|
||||
|
|
|
|||
Loading…
Reference in a new issue