Add code/cube3d/model/autoencoder/grid.py

code/cube3d/model/autoencoder/grid.py

@@ -0,0 +1,84 @@
+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