forked from openvdb/fvdb-core
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgrid_building.py
More file actions
150 lines (116 loc) · 5.58 KB
/
Copy pathgrid_building.py
File metadata and controls
150 lines (116 loc) · 5.58 KB
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
# Copyright Contributors to the OpenVDB Project
# SPDX-License-Identifier: Apache-2.0
#
from pathlib import Path
import fvdb.viz as fviz
import numpy as np
import point_cloud_utils as pcu
import polyscope as ps
import torch
from fvdb.types import DeviceIdentifier, resolve_device
from fvdb.utils.examples import load_car_1_mesh, load_car_2_mesh
import fvdb
from fvdb import JaggedTensor
voxel_size_1 = 0.02
voxel_size_2 = 0.03
def build_from_pointcloud(pcd_1: np.ndarray, pcd_2: np.ndarray):
# Assemble point clouds into JaggedTensor
pcd_jagged = JaggedTensor([torch.from_numpy(pcd_1).float().cuda(), torch.from_numpy(pcd_2).float().cuda()])
voxel_sizes = [[voxel_size_1, voxel_size_1, voxel_size_1], [voxel_size_2, voxel_size_2, voxel_size_2]]
# Method 1:
grid_a1 = fvdb.GridBatch.from_points(pcd_jagged, voxel_sizes=voxel_sizes, origins=[0.0] * 3)
# Visualization
gv_a1, ge_a1 = fviz.gridbatch_edge_network(grid_a1)
ps.remove_all_structures()
ps.register_point_cloud("pcd_1", pcd_1, enabled=True, radius=0.01)
ps.register_curve_network(
"grid_a1", gv_a1[0].jdata.cpu().numpy(), ge_a1[0].jdata.cpu().numpy(), enabled=True, radius=0.004
)
ps.show()
# Build grid from containing nearest voxels to the points
grid_b = fvdb.GridBatch.from_nearest_voxels_to_points(pcd_jagged, voxel_sizes=voxel_sizes, origins=[0.0] * 3)
# Visualization
gv_b, ge_b = fviz.gridbatch_edge_network(grid_b)
ps.remove_all_structures()
ps.register_point_cloud("pcd_1", pcd_1, enabled=True, radius=0.01)
ps.register_curve_network(
"grid_b1", gv_b[0].jdata.cpu().numpy(), ge_b[0].jdata.cpu().numpy(), enabled=True, radius=0.004
)
ps.show()
def build_from_coordinates(coords_1: np.ndarray, coords_2: np.ndarray):
coords_jagged = JaggedTensor([torch.from_numpy(coords_1).long().cuda(), torch.from_numpy(coords_2).long().cuda()])
voxel_sizes = [[voxel_size_1, voxel_size_1, voxel_size_1], [voxel_size_2, voxel_size_2, voxel_size_2]]
grid = fvdb.GridBatch.from_ijk(coords_jagged, voxel_sizes=voxel_sizes, origins=[0.0] * 3)
# Visualization
grid_mesh_1 = pcu.voxel_grid_geometry(
grid.ijk[0].jdata.cpu().numpy(), grid.voxel_sizes[0].cpu().numpy(), gap_fraction=0.1
)
grid_mesh_2 = pcu.voxel_grid_geometry(
grid.ijk[1].jdata.cpu().numpy(), grid.voxel_sizes[1].cpu().numpy(), gap_fraction=0.1
)
ps.remove_all_structures()
ps.register_surface_mesh("grid_1", grid_mesh_1[0], grid_mesh_1[1], enabled=True)
ps.register_surface_mesh("grid_2", grid_mesh_2[0], grid_mesh_2[1], enabled=True)
ps.show()
def build_from_mesh(mesh_1_vf, mesh_2_vf):
mesh_1_v, mesh_1_f = mesh_1_vf
mesh_2_v, mesh_2_f = mesh_2_vf
mesh_v_jagged = JaggedTensor([torch.from_numpy(mesh_1_v).float().cuda(), torch.from_numpy(mesh_2_v).float().cuda()])
mesh_f_jagged = JaggedTensor(
[
torch.from_numpy(mesh_1_f.astype(np.int64)).long().cuda(),
torch.from_numpy(mesh_2_f.astype(np.int64)).long().cuda(),
]
)
voxel_sizes = [[voxel_size_1, voxel_size_1, voxel_size_1], [voxel_size_2, voxel_size_2, voxel_size_2]]
grid = fvdb.GridBatch.from_mesh(mesh_v_jagged, mesh_f_jagged, voxel_sizes=voxel_sizes, origins=[0.0] * 3)
# Visualization
gv, ge = fviz.gridbatch_edge_network(grid)
ps.remove_all_structures()
ps.register_surface_mesh("mesh_1", mesh_1_v, mesh_1_f, enabled=True)
ps.register_curve_network(
"grid_1", gv[0].jdata.cpu().numpy(), ge[0].jdata.cpu().numpy(), enabled=True, radius=0.004
)
ps.register_surface_mesh("mesh_2", mesh_2_v, mesh_2_f, enabled=True)
ps.register_curve_network(
"grid_2", gv[1].jdata.cpu().numpy(), ge[1].jdata.cpu().numpy(), enabled=True, radius=0.004
)
ps.show()
def build_from_dense(device: DeviceIdentifier | None = None):
device = resolve_device(device)
batch_size = 2
dense_dims = [32, 32, 32]
voxel_sizes = [0.1] * 3
origins = [0.0] * 3
grid = fvdb.GridBatch.from_dense(
num_grids=batch_size, dense_dims=dense_dims, voxel_sizes=voxel_sizes, origins=origins, device=device
)
features_flat = torch.ones(grid.total_voxels, 16, device=device)
features = grid.jagged_like(features_flat)
# Visualization
grid_mesh = pcu.voxel_grid_geometry(
grid.ijk[0].jdata.cpu().numpy(), grid.voxel_sizes[0].cpu().numpy(), gap_fraction=0.1
)
ps.remove_all_structures()
ps.register_surface_mesh("grid_1", grid_mesh[0], grid_mesh[1], enabled=True)
ps.show()
if __name__ == "__main__":
ps.init()
ps.set_ground_plane_mode("shadow_only")
ps.set_navigation_style("free")
base_path = Path(__file__).parent.parent
mesh_1_v, mesh_1_f = load_car_1_mesh(mode="vf", device=torch.device("cpu"))
mesh_2_v, mesh_2_f = load_car_2_mesh(mode="vf", device=torch.device("cpu"))
mesh_1_v, mesh_1_f = mesh_1_v.numpy(), mesh_1_f.numpy().astype(np.int64)
mesh_2_v, mesh_2_f = mesh_2_v.numpy(), mesh_2_f.numpy().astype(np.int64)
mesh_2_v[:, 2] += 0.8
fi1, bc1 = pcu.sample_mesh_random(mesh_1_v, mesh_1_f, 10000)
fi2, bc2 = pcu.sample_mesh_random(mesh_2_v, mesh_2_f, 10000)
pcd_1 = pcu.interpolate_barycentric_coords(mesh_1_f, fi1, bc1, mesh_1_v)
pcd_2 = pcu.interpolate_barycentric_coords(mesh_2_f, fi2, bc2, mesh_2_v)
ijk_1 = np.unique(np.floor(pcd_1 / voxel_size_1).astype(np.int64), axis=0)
ijk_2 = np.unique(np.floor(pcd_2 / voxel_size_2).astype(np.int64), axis=0)
build_from_pointcloud(pcd_1, pcd_2)
build_from_mesh((mesh_1_v, mesh_1_f), (mesh_2_v, mesh_2_f))
build_from_coordinates(ijk_1, ijk_2)
build_from_dense()