| from typing import Literal, Union | |
| import numpy as np | |
| import torch | |
| import warp as wp | |
| def generate_dense_grid_points( | |
| bbox_min: np.ndarray, | |
| bbox_max: np.ndarray, | |
| resolution_base: float, | |
| indexing: Literal["xy", "ij"] = "ij", | |
| ) -> tuple[np.ndarray, list[int], np.ndarray]: | |
| """ | |
| Generate a dense grid of points within a bounding box. | |
| Parameters: | |
| bbox_min (np.ndarray): The minimum coordinates of the bounding box (3D). | |
| bbox_max (np.ndarray): The maximum coordinates of the bounding box (3D). | |
| resolution_base (float): The base resolution for the grid. The number of cells along each axis will be 2^resolution_base. | |
| indexing (Literal["xy", "ij"], optional): The indexing convention for the grid. "xy" for Cartesian indexing, "ij" for matrix indexing. Default is "ij". | |
| Returns: | |
| tuple: A tuple containing: | |
| - xyz (np.ndarray): A 2D array of shape (N, 3) where N is the total number of grid points. Each row represents the (x, y, z) coordinates of a grid point. | |
| - grid_size (list): A list of three integers representing the number of grid points along each axis. | |
| - length (np.ndarray): The length of the bounding box along each axis. | |
| """ | |
| length = bbox_max - bbox_min | |
| num_cells = np.exp2(resolution_base) | |
| x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32) | |
| y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32) | |
| z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32) | |
| [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) | |
| xyz = np.stack((xs, ys, zs), axis=-1) | |
| xyz = xyz.reshape(-1, 3) | |
| grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1] | |
| return xyz, grid_size, length | |
| def marching_cubes_with_warp( | |
| grid_logits: torch.Tensor, | |
| level: float, | |
| device: Union[str, torch.device] = "cuda", | |
| max_verts: int = 3_000_000, | |
| max_tris: int = 3_000_000, | |
| ) -> tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Perform the marching cubes algorithm on a 3D grid with warp support. | |
| Args: | |
| grid_logits (torch.Tensor): A 3D tensor containing the grid logits. | |
| level (float): The threshold level for the isosurface. | |
| device (Union[str, torch.device], optional): The device to perform the computation on. Defaults to "cuda". | |
| max_verts (int, optional): The maximum number of vertices. Defaults to 3,000,000. | |
| max_tris (int, optional): The maximum number of triangles. Defaults to 3,000,000. | |
| Returns: | |
| Tuple[np.ndarray, np.ndarray]: A tuple containing the vertices and faces of the isosurface. | |
| """ | |
| if isinstance(device, torch.device): | |
| device = str(device) | |
| assert grid_logits.ndim == 3 | |
| if "cuda" in device: | |
| assert wp.is_cuda_available() | |
| else: | |
| raise ValueError( | |
| f"Device {device} is not supported for marching_cubes_with_warp" | |
| ) | |
| dim = grid_logits.shape[0] | |
| field = wp.from_torch(grid_logits) | |
| iso = wp.MarchingCubes( | |
| nx=dim, | |
| ny=dim, | |
| nz=dim, | |
| max_verts=int(max_verts), | |
| max_tris=int(max_tris), | |
| device=device, | |
| ) | |
| iso.surface(field=field, threshold=level) | |
| vertices = iso.verts.numpy() | |
| faces = iso.indices.numpy().reshape(-1, 3) | |
| return vertices, faces | |