tohid4n/PartCrafter
0
1from src.utils.typing_utils import *2 3import numpy as np4import torch5import torch.nn as nn6import scipy.ndimage7from skimage import measure8from einops import repeat9from diso import DiffDMC10import torch.nn.functional as F11 12def generate_dense_grid_points(13 bbox_min: np.ndarray, bbox_max: np.ndarray, octree_depth: int, indexing: str = "ij"14):15 length = bbox_max - bbox_min16 num_cells = np.exp2(octree_depth)17 x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)18 y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)19 z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)20 [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)21 xyz = np.stack((xs, ys, zs), axis=-1)22 xyz = xyz.reshape(-1, 3)23 grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]24 25 return xyz, grid_size, length26 27def generate_dense_grid_points_gpu(28 bbox_min: torch.Tensor, 29 bbox_max: torch.Tensor,30 octree_depth: int,31 indexing: str = "ij", 32 dtype: torch.dtype = torch.float1633):34 length = bbox_max - bbox_min35 num_cells = 2 ** octree_depth36 device = bbox_min.device37 38 x = torch.linspace(bbox_min[0], bbox_max[0], int(num_cells), dtype=dtype, device=device)39 y = torch.linspace(bbox_min[1], bbox_max[1], int(num_cells), dtype=dtype, device=device)40 z = torch.linspace(bbox_min[2], bbox_max[2], int(num_cells), dtype=dtype, device=device)41 42 xs, ys, zs = torch.meshgrid(x, y, z, indexing=indexing)43 xyz = torch.stack((xs, ys, zs), dim=-1)44 xyz = xyz.view(-1, 3)45 grid_size = [int(num_cells), int(num_cells), int(num_cells)]46 47 return xyz, grid_size, length48 49def find_mesh_grid_coordinates_fast_gpu(50 occupancy_grid, 51 n_limits=-152):53 core_grid = occupancy_grid[1:-1, 1:-1, 1:-1]54 occupied = core_grid > 055 56 neighbors_unoccupied = (57 (occupancy_grid[:-2, :-2, :-2] < 0)58 | (occupancy_grid[:-2, :-2, 1:-1] < 0)59 | (occupancy_grid[:-2, :-2, 2:] < 0) # x-1, y-1, z-1/0/160 | (occupancy_grid[:-2, 1:-1, :-2] < 0)61 | (occupancy_grid[:-2, 1:-1, 1:-1] < 0)62 | (occupancy_grid[:-2, 1:-1, 2:] < 0) # x-1, y0, z-1/0/163 | (occupancy_grid[:-2, 2:, :-2] < 0)64 | (occupancy_grid[:-2, 2:, 1:-1] < 0)65 | (occupancy_grid[:-2, 2:, 2:] < 0) # x-1, y+1, z-1/0/166 | (occupancy_grid[1:-1, :-2, :-2] < 0)67 | (occupancy_grid[1:-1, :-2, 1:-1] < 0)68 | (occupancy_grid[1:-1, :-2, 2:] < 0) # x0, y-1, z-1/0/169 | (occupancy_grid[1:-1, 1:-1, :-2] < 0)70 | (occupancy_grid[1:-1, 1:-1, 2:] < 0) # x0, y0, z-1/171 | (occupancy_grid[1:-1, 2:, :-2] < 0)72 | (occupancy_grid[1:-1, 2:, 1:-1] < 0)73 | (occupancy_grid[1:-1, 2:, 2:] < 0) # x0, y+1, z-1/0/174 | (occupancy_grid[2:, :-2, :-2] < 0)75 | (occupancy_grid[2:, :-2, 1:-1] < 0)76 | (occupancy_grid[2:, :-2, 2:] < 0) # x+1, y-1, z-1/0/177 | (occupancy_grid[2:, 1:-1, :-2] < 0)78 | (occupancy_grid[2:, 1:-1, 1:-1] < 0)79 | (occupancy_grid[2:, 1:-1, 2:] < 0) # x+1, y0, z-1/0/180 | (occupancy_grid[2:, 2:, :-2] < 0)81 | (occupancy_grid[2:, 2:, 1:-1] < 0)82 | (occupancy_grid[2:, 2:, 2:] < 0) # x+1, y+1, z-1/0/183 )84 core_mesh_coords = torch.nonzero(occupied & neighbors_unoccupied, as_tuple=False) + 185 86 if n_limits != -1 and core_mesh_coords.shape[0] > n_limits:87 print(f"core mesh coords {core_mesh_coords.shape[0]} is too large, limited to {n_limits}")88 ind = np.random.choice(core_mesh_coords.shape[0], n_limits, True)89 core_mesh_coords = core_mesh_coords[ind]90 91 return core_mesh_coords92 93def find_candidates_band(94 occupancy_grid: torch.Tensor, 95 band_threshold: float, 96 n_limits: int = -197) -> torch.Tensor:98 """99 Returns the coordinates of all voxels in the occupancy_grid where |value| < band_threshold.100 101 Args:102 occupancy_grid (torch.Tensor): A 3D tensor of SDF values.103 band_threshold (float): The threshold below which |SDF| must be to include the voxel.104 n_limits (int): Maximum number of points to return (-1 for no limit)105 106 Returns:107 torch.Tensor: A 2D tensor of coordinates (N x 3) where each row is [x, y, z].108 """109 core_grid = occupancy_grid[1:-1, 1:-1, 1:-1] 110 # logits to sdf111 core_grid = torch.sigmoid(core_grid) * 2 - 1 112 # Create a boolean mask for all cells in the band113 in_band = torch.abs(core_grid) < band_threshold114 115 # Get coordinates of all voxels in the band116 core_mesh_coords = torch.nonzero(in_band, as_tuple=False) + 1117 118 if n_limits != -1 and core_mesh_coords.shape[0] > n_limits:119 print(f"core mesh coords {core_mesh_coords.shape[0]} is too large, limited to {n_limits}")120 ind = np.random.choice(core_mesh_coords.shape[0], n_limits, True)121 core_mesh_coords = core_mesh_coords[ind]122 123 return core_mesh_coords 124 125def expand_edge_region_fast(edge_coords, grid_size, dtype):126 expanded_tensor = torch.zeros(grid_size, grid_size, grid_size, device='cuda', dtype=dtype, requires_grad=False)127 expanded_tensor[edge_coords[:, 0], edge_coords[:, 1], edge_coords[:, 2]] = 1128 if grid_size < 512:129 kernel_size = 5130 pooled_tensor = torch.nn.functional.max_pool3d(expanded_tensor.unsqueeze(0).unsqueeze(0), kernel_size=kernel_size, stride=1, padding=2).squeeze()131 else:132 kernel_size = 3133 pooled_tensor = torch.nn.functional.max_pool3d(expanded_tensor.unsqueeze(0).unsqueeze(0), kernel_size=kernel_size, stride=1, padding=1).squeeze()134 expanded_coords_low_res = torch.nonzero(pooled_tensor, as_tuple=False).to(torch.int16)135 136 expanded_coords_high_res = torch.stack([137 torch.cat((expanded_coords_low_res[:, 0] * 2, expanded_coords_low_res[:, 0] * 2, expanded_coords_low_res[:, 0] * 2, expanded_coords_low_res[:, 0] * 2, expanded_coords_low_res[:, 0] * 2 + 1, expanded_coords_low_res[:, 0] * 2 + 1, expanded_coords_low_res[:, 0] * 2 + 1, expanded_coords_low_res[:, 0] * 2 + 1)),138 torch.cat((expanded_coords_low_res[:, 1] * 2, expanded_coords_low_res[:, 1] * 2, expanded_coords_low_res[:, 1] * 2+1, expanded_coords_low_res[:, 1] * 2 + 1, expanded_coords_low_res[:, 1] * 2, expanded_coords_low_res[:, 1] * 2, expanded_coords_low_res[:, 1] * 2 + 1, expanded_coords_low_res[:, 1] * 2 + 1)),139 torch.cat((expanded_coords_low_res[:, 2] * 2, expanded_coords_low_res[:, 2] * 2+1, expanded_coords_low_res[:, 2] * 2, expanded_coords_low_res[:, 2] * 2 + 1, expanded_coords_low_res[:, 2] * 2, expanded_coords_low_res[:, 2] * 2+1, expanded_coords_low_res[:, 2] * 2, expanded_coords_low_res[:, 2] * 2 + 1))140 ], dim=1)141 142 return expanded_coords_high_res143 144def zoom_block(block, scale_factor, order=3):145 block = block.astype(np.float32)146 return scipy.ndimage.zoom(block, scale_factor, order=order)147 148def parallel_zoom(occupancy_grid, scale_factor):149 result = torch.nn.functional.interpolate(occupancy_grid.unsqueeze(0).unsqueeze(0), scale_factor=scale_factor)150 return result.squeeze(0).squeeze(0)151 152 153@torch.no_grad()154def hierarchical_extract_geometry(155 geometric_func: Callable,156 device: torch.device,157 dtype: torch.dtype,158 bounds: Union[Tuple[float], List[float], float] = (-1.25, -1.25, -1.25, 1.25, 1.25, 1.25),159 dense_octree_depth: int = 8,160 hierarchical_octree_depth: int = 9, 161 max_num_expanded_coords: int = 1e8, 162 verbose: bool = False,163):164 """165 Args:166 geometric_func:167 device:168 bounds:169 dense_octree_depth:170 hierarchical_octree_depth:171 Returns:172 """173 if isinstance(bounds, float):174 bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]175 176 bbox_min = torch.tensor(bounds[0:3]).to(device)177 bbox_max = torch.tensor(bounds[3:6]).to(device)178 bbox_size = bbox_max - bbox_min179 180 xyz_samples, grid_size, length = generate_dense_grid_points_gpu(181 bbox_min=bbox_min,182 bbox_max=bbox_max,183 octree_depth=dense_octree_depth,184 indexing="ij",185 dtype=dtype186 )187 188 if verbose:189 print(f'step 1 query num: {xyz_samples.shape[0]}')190 grid_logits = geometric_func(xyz_samples.unsqueeze(0)).to(dtype).view(grid_size[0], grid_size[1], grid_size[2])191 # print(f'step 1 grid_logits shape: {grid_logits.shape}')192 for i in range(hierarchical_octree_depth - dense_octree_depth):193 curr_octree_depth = dense_octree_depth + i + 1194 # upsample195 grid_size = 2**curr_octree_depth196 normalize_offset = grid_size / 2197 high_res_occupancy = parallel_zoom(grid_logits, 2).to(dtype)198 199 band_threshold = 1.0200 edge_coords = find_candidates_band(grid_logits, band_threshold)201 expanded_coords = expand_edge_region_fast(edge_coords, grid_size=int(grid_size/2), dtype=dtype).to(dtype)202 if verbose:203 print(f'step {i+2} query num: {len(expanded_coords)}')204 if max_num_expanded_coords > 0 and len(expanded_coords) > max_num_expanded_coords:205 raise ValueError(f"expanded_coords is too large, {len(expanded_coords)} > {max_num_expanded_coords}")206 expanded_coords_norm = (expanded_coords - normalize_offset) * (abs(bounds[0]) / normalize_offset)207 208 all_logits = None209 210 all_logits = geometric_func(expanded_coords_norm.unsqueeze(0)).to(dtype)211 all_logits = torch.cat([expanded_coords_norm, all_logits[0]], dim=1)212 # print("all logits shape = ", all_logits.shape)213 214 indices = all_logits[..., :3]215 indices = indices * (normalize_offset / abs(bounds[0])) + normalize_offset216 indices = indices.type(torch.IntTensor)217 values = all_logits[:, 3]218 # breakpoint()219 high_res_occupancy[indices[:, 0], indices[:, 1], indices[:, 2]] = values220 grid_logits = high_res_occupancy221 # torch.cuda.empty_cache()222 223 if verbose:224 print("final grids shape = ", grid_logits.shape)225 vertices, faces, normals, _ = measure.marching_cubes(grid_logits.float().cpu().numpy(), 0, method="lewiner")226 vertices = vertices / (2**hierarchical_octree_depth) * bbox_size.cpu().numpy() + bbox_min.cpu().numpy()227 mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))228 229 return mesh_v_f230 231def extract_near_surface_volume_fn(input_tensor: torch.Tensor, alpha: float):232 """233 Args:234 input_tensor: shape [D, D, D], torch.float16235 alpha: isosurface offset236 Returns:237 mask: shape [D, D, D], torch.int32238 """239 device = input_tensor.device240 D = input_tensor.shape[0]241 signed_val = 0.0242 243 # add isosurface offset and exclude invalid value244 val = input_tensor + alpha245 valid_mask = val > -9000246 247 # obtain neighbors248 def get_neighbor(t, shift, axis):249 if shift == 0:250 return t.clone()251 252 pad_dims = [0, 0, 0, 0, 0, 0] # [x_front，x_back，y_front，y_back，z_front，z_back]253 254 if axis == 0: # x axis255 pad_idx = 0 if shift > 0 else 1256 pad_dims[pad_idx] = abs(shift)257 elif axis == 1: # y axis258 pad_idx = 2 if shift > 0 else 3259 pad_dims[pad_idx] = abs(shift)260 elif axis == 2: # z axis261 pad_idx = 4 if shift > 0 else 5262 pad_dims[pad_idx] = abs(shift)263 264 # Apply padding with replication at boundaries265 padded = F.pad(t.unsqueeze(0).unsqueeze(0), pad_dims[::-1], mode='replicate')266 267 # Create dynamic slicing indices268 slice_dims = [slice(None)] * 3269 if axis == 0: # x axis270 if shift > 0:271 slice_dims[0] = slice(shift, None)272 else:273 slice_dims[0] = slice(None, shift)274 elif axis == 1: # y axis275 if shift > 0:276 slice_dims[1] = slice(shift, None)277 else:278 slice_dims[1] = slice(None, shift)279 elif axis == 2: # z axis280 if shift > 0:281 slice_dims[2] = slice(shift, None)282 else:283 slice_dims[2] = slice(None, shift)284 285 # Apply slicing and restore dimensions286 padded = padded.squeeze(0).squeeze(0)287 sliced = padded[slice_dims]288 return sliced289 290 # Get neighbors in all directions291 left = get_neighbor(val, 1, axis=0) # x axis292 right = get_neighbor(val, -1, axis=0)293 back = get_neighbor(val, 1, axis=1) # y axis294 front = get_neighbor(val, -1, axis=1)295 down = get_neighbor(val, 1, axis=2) # z axis296 up = get_neighbor(val, -1, axis=2)297 298 # Handle invalid boundary values299 def safe_where(neighbor):300 return torch.where(neighbor > -9000, neighbor, val)301 302 left = safe_where(left)303 right = safe_where(right)304 back = safe_where(back)305 front = safe_where(front)306 down = safe_where(down)307 up = safe_where(up)308 309 # Calculate sign consistency310 sign = torch.sign(val.to(torch.float32))311 neighbors_sign = torch.stack([312 torch.sign(left.to(torch.float32)),313 torch.sign(right.to(torch.float32)),314 torch.sign(back.to(torch.float32)),315 torch.sign(front.to(torch.float32)),316 torch.sign(down.to(torch.float32)),317 torch.sign(up.to(torch.float32))318 ], dim=0)319 320 # Check if all signs are consistent321 same_sign = torch.all(neighbors_sign == sign, dim=0)322 323 # Generate final mask324 mask = (~same_sign).to(torch.int32)325 return mask * valid_mask.to(torch.int32)326 327 328def generate_dense_grid_points_2(329 bbox_min: np.ndarray,330 bbox_max: np.ndarray,331 octree_resolution: int,332 indexing: str = "ij",333):334 length = bbox_max - bbox_min335 num_cells = octree_resolution336 337 x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)338 y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)339 z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)340 [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)341 xyz = np.stack((xs, ys, zs), axis=-1)342 grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]343 344 return xyz, grid_size, length345 346@torch.no_grad()347def flash_extract_geometry(348 latents: torch.FloatTensor,349 vae: Callable,350 bounds: Union[Tuple[float], List[float], float] = 1.01,351 num_chunks: int = 10000,352 mc_level: float = 0.0,353 octree_depth: int = 9,354 min_resolution: int = 63,355 mini_grid_num: int = 4,356 **kwargs,357):358 geo_decoder = vae.decoder359 device = latents.device360 dtype = latents.dtype361 # resolution to depth362 octree_resolution = 2 ** octree_depth363 resolutions = []364 if octree_resolution < min_resolution:365 resolutions.append(octree_resolution)366 while octree_resolution >= min_resolution:367 resolutions.append(octree_resolution)368 octree_resolution = octree_resolution // 2369 resolutions.reverse()370 resolutions[0] = round(resolutions[0] / mini_grid_num) * mini_grid_num - 1371 for i, resolution in enumerate(resolutions[1:]):372 resolutions[i + 1] = resolutions[0] * 2 ** (i + 1)373 374 375 # 1. generate query points376 if isinstance(bounds, float):377 bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]378 bbox_min = np.array(bounds[0:3])379 bbox_max = np.array(bounds[3:6])380 bbox_size = bbox_max - bbox_min381 382 xyz_samples, grid_size, length = generate_dense_grid_points_2(383 bbox_min=bbox_min,384 bbox_max=bbox_max,385 octree_resolution=resolutions[0],386 indexing="ij"387 )388 389 dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype)390 dilate.weight = torch.nn.Parameter(torch.ones(dilate.weight.shape, dtype=dtype, device=device))391 392 grid_size = np.array(grid_size)393 394 # 2. latents to 3d volume395 xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype)396 batch_size = latents.shape[0]397 mini_grid_size = xyz_samples.shape[0] // mini_grid_num398 xyz_samples = xyz_samples.view(399 mini_grid_num, mini_grid_size,400 mini_grid_num, mini_grid_size,401 mini_grid_num, mini_grid_size, 3402 ).permute(403 0, 2, 4, 1, 3, 5, 6404 ).reshape(405 -1, mini_grid_size * mini_grid_size * mini_grid_size, 3406 )407 batch_logits = []408 num_batchs = max(num_chunks // xyz_samples.shape[1], 1)409 for start in range(0, xyz_samples.shape[0], num_batchs):410 queries = xyz_samples[start: start + num_batchs, :]411 batch = queries.shape[0]412 batch_latents = repeat(latents.squeeze(0), "p c -> b p c", b=batch)413 # geo_decoder.set_topk(True)414 geo_decoder.set_topk(False)415 logits = vae.decode(batch_latents, queries).sample416 batch_logits.append(logits)417 grid_logits = torch.cat(batch_logits, dim=0).reshape(418 mini_grid_num, mini_grid_num, mini_grid_num,419 mini_grid_size, mini_grid_size,420 mini_grid_size421 ).permute(0, 3, 1, 4, 2, 5).contiguous().view(422 (batch_size, grid_size[0], grid_size[1], grid_size[2])423 )424 425 for octree_depth_now in resolutions[1:]:426 grid_size = np.array([octree_depth_now + 1] * 3)427 resolution = bbox_size / octree_depth_now428 next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device)429 next_logits = torch.full(next_index.shape, -10000., dtype=dtype, device=device)430 curr_points = extract_near_surface_volume_fn(grid_logits.squeeze(0), mc_level)431 curr_points += grid_logits.squeeze(0).abs() < 0.95432 433 if octree_depth_now == resolutions[-1]:434 expand_num = 0435 else:436 expand_num = 1437 for i in range(expand_num):438 curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0)439 curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0)440 (cidx_x, cidx_y, cidx_z) = torch.where(curr_points > 0)441 442 next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1443 for i in range(2 - expand_num):444 next_index = dilate(next_index.unsqueeze(0)).squeeze(0)445 nidx = torch.where(next_index > 0)446 447 next_points = torch.stack(nidx, dim=1)448 next_points = (next_points * torch.tensor(resolution, dtype=torch.float32, device=device) +449 torch.tensor(bbox_min, dtype=torch.float32, device=device))450 451 query_grid_num = 6452 min_val = next_points.min(axis=0).values453 max_val = next_points.max(axis=0).values454 vol_queries_index = (next_points - min_val) / (max_val - min_val) * (query_grid_num - 0.001)455 index = torch.floor(vol_queries_index).long()456 index = index[..., 0] * (query_grid_num ** 2) + index[..., 1] * query_grid_num + index[..., 2]457 index = index.sort()458 next_points = next_points[index.indices].unsqueeze(0).contiguous()459 unique_values = torch.unique(index.values, return_counts=True)460 grid_logits = torch.zeros((next_points.shape[1]), dtype=latents.dtype, device=latents.device)461 input_grid = [[], []]462 logits_grid_list = []463 start_num = 0464 sum_num = 0465 for grid_index, count in zip(unique_values[0].cpu().tolist(), unique_values[1].cpu().tolist()):466 if sum_num + count < num_chunks or sum_num == 0:467 sum_num += count468 input_grid[0].append(grid_index)469 input_grid[1].append(count)470 else:471 # geo_decoder.set_topk(input_grid)472 geo_decoder.set_topk(False)473 logits_grid = vae.decode(latents,next_points[:, start_num:start_num + sum_num]).sample474 start_num = start_num + sum_num475 logits_grid_list.append(logits_grid)476 input_grid = [[grid_index], [count]]477 sum_num = count478 if sum_num > 0:479 # geo_decoder.set_topk(input_grid)480 geo_decoder.set_topk(False)481 logits_grid = vae.decode(latents,next_points[:, start_num:start_num + sum_num]).sample482 logits_grid_list.append(logits_grid)483 logits_grid = torch.cat(logits_grid_list, dim=1)484 grid_logits[index.indices] = logits_grid.squeeze(0).squeeze(-1)485 next_logits[nidx] = grid_logits486 grid_logits = next_logits.unsqueeze(0)487 488 grid_logits[grid_logits == -10000.] = float('nan')489 torch.cuda.empty_cache()490 mesh_v_f = []491 grid_logits = grid_logits[0]492 try:493 print("final grids shape = ", grid_logits.shape)494 dmc = DiffDMC(dtype=torch.float32).to(grid_logits.device)495 sdf = -grid_logits / octree_resolution496 sdf = sdf.to(torch.float32).contiguous()497 vertices, faces = dmc(sdf, deform=None, return_quads=False, normalize=False)498 vertices = vertices.detach().cpu().numpy()499 faces = faces.detach().cpu().numpy()[:, ::-1] 500 vertices = vertices / (2 ** octree_depth) * bbox_size + bbox_min501 mesh_v_f = (vertices.astype(np.float32), np.ascontiguousarray(faces))502 except Exception as e:503 print(e)504 torch.cuda.empty_cache()505 mesh_v_f = (None, None)506 507 return [mesh_v_f]508