conquer3d 0.2.9__tar.gz → 0.3.0__tar.gz

{conquer3d-0.2.9/conquer3d.egg-info → conquer3d-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.9
+Version: 0.3.0
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT
@@ -51,8 +51,13 @@ conda activate geocutool
 
 conda install nvidia::cuda-toolkit==12.8.2 -y
 
+pip install setuptools wheel ninja
 pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
 
+pip install git+https://github.com/mit-han-lab/torchsparse.git
+
+pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.8.0_cu128.html
+
 pip install pybind11-stubgen
 pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
 

{conquer3d-0.2.9 → conquer3d-0.3.0}/README.md

@@ -23,8 +23,13 @@ conda activate geocutool
 
 conda install nvidia::cuda-toolkit==12.8.2 -y
 
+pip install setuptools wheel ninja
 pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
 
+pip install git+https://github.com/mit-han-lab/torchsparse.git
+
+pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.8.0_cu128.html
+
 pip install pybind11-stubgen
 pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
 

{conquer3d-0.2.9 → conquer3d-0.3.0}/conquer3d/csrc/binds/data_structure/grid.cpp

@@ -121,9 +121,30 @@ torch::Tensor compute_grid_normal(torch::Tensor sdf, torch::Tensor grid_vertices
     return torch::nn::functional::normalize(normals, torch::nn::functional::NormalizeFuncOptions().p(2).dim(1));
 }
 
+torch::Tensor compute_active_voxels(torch::Tensor voxels, torch::Tensor sdf, float iso) {
+    TORCH_CHECK(voxels.is_cuda(), "voxels must be a CUDA tensor.");
+    TORCH_CHECK(sdf.is_cuda(), "sdf must be a CUDA tensor.");
+    TORCH_CHECK(voxels.dtype() == torch::kInt32, "voxels must be int32.");
+    TORCH_CHECK(sdf.dtype() == torch::kFloat32, "sdf must be float32.");
+    
+    auto voxels_flat = voxels.flatten().to(torch::kInt64);
+    auto voxel_sdfs = sdf.index_select(0, voxels_flat).view({-1, 8});
+    
+    auto below_iso = voxel_sdfs < iso;
+    auto any_below = below_iso.any(1);
+    auto any_above = (~below_iso).any(1);
+    
+    auto active_mask = any_below & any_above;
+    auto active_indices = torch::nonzero(active_mask).squeeze(1);
+    
+    return active_indices;
+}
+
 void bind_ds_grid(py::module_& m) {
     m.def("create_voxel_grid", &create_voxel_grid, "Creates a structured 3D voxel grid.",
           py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("device_str") = "cuda");
     m.def("compute_grid_normal", &compute_grid_normal, "Computes surface normals for a grid.",
           py::arg("sdf"), py::arg("grid_vertices"), py::arg("idx_grids"), py::arg("res"));
+    m.def("compute_active_voxels", &compute_active_voxels, "Computes active voxels intersecting the surface",
+          py::arg("voxels"), py::arg("sdf"), py::arg("iso"));
 }

conquer3d-0.3.0/conquer3d/csrc/binds/ops/mc.cpp

@@ -0,0 +1,118 @@
+#include <torch/extension.h>
+#include <pybind11/pybind11.h>
+#include "../../ops/mc.h"
+#include "../../check.h"
+
+namespace py = pybind11;
+
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Tensor>, std::optional<torch::Tensor>> marching_cubes_wrapper(
+    torch::Tensor grid_vertices,
+    torch::Tensor voxels,
+    torch::Tensor voxel_values,
+    std::optional<torch::Tensor> grid_normals,
+    std::optional<torch::Tensor> grid_colors,
+    float iso,
+    bool return_unique_edges
+) {
+    CHECK_INPUT(grid_vertices);
+    CHECK_INPUT(voxels);
+    CHECK_INPUT(voxel_values);
+
+    uint32_t num_voxels = voxels.size(0);
+
+    const float3* __restrict__ p_grid_vertices = (float3*)grid_vertices.data_ptr<float>();
+    const uint32_t* __restrict__ p_voxels = (uint32_t*)voxels.data_ptr<int32_t>();
+    const float* __restrict__ p_voxel_values = voxel_values.data_ptr<float>();
+
+    const float3* __restrict__ p_grid_normals = nullptr;
+    if (grid_normals.has_value()) {
+        CHECK_INPUT(grid_normals.value());
+        p_grid_normals = (float3*)grid_normals.value().data_ptr<float>();
+    }
+
+    const float3* __restrict__ p_grid_colors = nullptr;
+    if (grid_colors.has_value()) {
+        CHECK_INPUT(grid_colors.value());
+        p_grid_colors = (float3*)grid_colors.value().data_ptr<float>();
+    }
+
+    return mc::marching_cubes(
+        num_voxels,
+        p_grid_vertices,
+        p_voxels,
+        p_voxel_values,
+        p_grid_normals,
+        p_grid_colors,
+        iso,
+        grid_vertices.options(),
+        voxels.options(),
+        return_unique_edges
+    );
+}
+
+void marching_cubes_backward_wrapper(
+    torch::Tensor unique_edges,
+    torch::Tensor grid_vertices,
+    std::optional<torch::Tensor> grid_colors,
+    torch::Tensor values,
+    torch::Tensor adj_verts,
+    std::optional<torch::Tensor> adj_colors,
+    torch::Tensor adj_values,
+    std::optional<torch::Tensor> adj_grid_colors,
+    float iso
+) {
+    CHECK_INPUT(unique_edges);
+    CHECK_INPUT(grid_vertices);
+    CHECK_INPUT(values);
+    CHECK_INPUT(adj_verts);
+    CHECK_INPUT(adj_values);
+
+    uint32_t n_verts = unique_edges.size(0);
+
+    const Edge* p_unique_edges = (const Edge*)unique_edges.data_ptr<int32_t>();
+    const float3* p_grid_vertices = (const float3*)grid_vertices.data_ptr<float>();
+    const float* p_values = values.data_ptr<float>();
+    const float3* p_adj_verts = (const float3*)adj_verts.data_ptr<float>();
+    float* p_adj_values = adj_values.data_ptr<float>();
+
+    const float3* p_grid_colors = nullptr;
+    if (grid_colors.has_value()) {
+        CHECK_INPUT(grid_colors.value());
+        p_grid_colors = (const float3*)grid_colors.value().data_ptr<float>();
+    }
+
+    const float3* p_adj_colors = nullptr;
+    if (adj_colors.has_value()) {
+        CHECK_INPUT(adj_colors.value());
+        p_adj_colors = (const float3*)adj_colors.value().data_ptr<float>();
+    }
+
+    float3* p_adj_grid_colors = nullptr;
+    if (adj_grid_colors.has_value()) {
+        CHECK_INPUT(adj_grid_colors.value());
+        p_adj_grid_colors = (float3*)adj_grid_colors.value().data_ptr<float>();
+    }
+
+    mc::backward(
+        n_verts,
+        p_unique_edges,
+        p_grid_vertices,
+        p_grid_colors,
+        p_values,
+        p_adj_verts,
+        p_adj_colors,
+        p_adj_values,
+        p_adj_grid_colors,
+        iso
+    );
+}
+
+void bind_ops_mc(py::module_& m) {
+    m.def("marching_cubes", &marching_cubes_wrapper, "Marching Cubes",
+          py::arg("grid_vertices"), py::arg("voxels"), py::arg("voxel_values"),
+          py::arg("grid_normals") = std::nullopt, py::arg("grid_colors") = std::nullopt, py::arg("iso") = 0.0f, py::arg("return_unique_edges") = false);
+
+    m.def("marching_cubes_backward", &marching_cubes_backward_wrapper, "Marching Cubes Backward",
+          py::arg("unique_edges"), py::arg("grid_vertices"), py::arg("grid_colors"), py::arg("values"),
+          py::arg("adj_verts"), py::arg("adj_colors"), py::arg("adj_values"), py::arg("adj_grid_colors"), py::arg("iso"));
+}

{conquer3d-0.2.9 → conquer3d-0.3.0}/conquer3d/csrc/ops/mc.cu

@@ -176,9 +176,11 @@ namespace mc
         const float3* grid_vertices,
         const float* values,
         const float3* grid_normals,
+        const float3* grid_colors,
         const float iso,
         float3* out_verts,
-        float3* out_normals
+        float3* out_normals,
+        float3* out_colors
     )
     {
         uint32_t v_idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -199,24 +201,31 @@ namespace mc
         // 3. Interpolate (Differentiable formula)
         float3 p;
         float3 n;
+        float3 c;
         bool has_normals = grid_normals != nullptr;
+        bool has_colors = grid_colors != nullptr;
         float3 n0 = has_normals ? grid_normals[v0_idx] : make_float3(0, 0, 0);
         float3 n1 = has_normals ? grid_normals[v1_idx] : make_float3(0, 0, 0);
+        float3 c0 = has_colors ? grid_colors[v0_idx] : make_float3(0, 0, 0);
+        float3 c1 = has_colors ? grid_colors[v1_idx] : make_float3(0, 0, 0);
 
         if (fabsf(iso - val0) < EPS)
         {
             p = p0;
             if (has_normals) n = n0;
+            if (has_colors) c = c0;
         }
         else if (fabsf(iso - val1) < EPS)
         {
             p = p1;
             if (has_normals) n = n1;
+            if (has_colors) c = c1;
         }
         else if (fabsf(val0 - val1) < EPS)
         {
             p = p0;
             if (has_normals) n = n0;
+            if (has_colors) c = c0;
         }
         else
         {
@@ -228,12 +237,18 @@ namespace mc
                 n = n0 + (n1 - n0) * t;
                 n = maths::normalize(n);
             }
+            if (has_colors) {
+                c = c0 + (c1 - c0) * t;
+            }
         }
 
         out_verts[v_idx] = p;
         if (has_normals) {
             out_normals[v_idx] = n;
         }
+        if (has_colors) {
+            out_colors[v_idx] = c;
+        }
     }
 
     void compute_active_voxels(
@@ -357,16 +372,18 @@ namespace mc
         const float3* grid_vertices,
         const float* values,
         const float3* grid_normals,
+        const float3* grid_colors,
         const float iso,
         float3* out_verts,
-        float3* out_normals
+        float3* out_normals,
+        float3* out_colors
     )
     {
         if (num_unique_edges == 0) return;
         int block_size = NTHREADS;
         int grid_size = (num_unique_edges + block_size - 1) / block_size;
         interpolate_vertices_kernel<<<grid_size, block_size>>>(
-            num_unique_edges, unique_edges, grid_vertices, values, grid_normals, iso, out_verts, out_normals);
+            num_unique_edges, unique_edges, grid_vertices, values, grid_normals, grid_colors, iso, out_verts, out_normals, out_colors);
     }
 
     void compute_number_triangles(
@@ -435,15 +452,17 @@ namespace mc
             num_active_voxels, used_voxel_codes, voxel_edge_to_vert_idx, voxel_triangle_prefix_sums, out_triangles);
     }
 
-    std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> marching_cubes(
+    std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Tensor>, std::optional<torch::Tensor>> marching_cubes(
         const uint32_t num_voxels,
         const float3* __restrict__ grid_vertices,
         const uint32_t* __restrict__ voxels,
         const float* __restrict__ voxel_values,
         const float3* __restrict__ grid_normals,
+        const float3* __restrict__ grid_colors,
         const float iso,
         torch::TensorOptions vert_options,
-        torch::TensorOptions tri_options
+        torch::TensorOptions tri_options,
+        bool return_unique_edges
     )
     {
         uint8_t *__restrict__ voxel_codes;
@@ -457,11 +476,15 @@ namespace mc
         if (num_active_voxels == 0) {
             CHECK_CUDA_INTERNAL(cudaFree(voxel_codes));
             std::optional<torch::Tensor> out_n = std::nullopt;
+            std::optional<torch::Tensor> out_c = std::nullopt;
             if (grid_normals != nullptr) out_n = torch::empty({0, 3}, vert_options);
+            if (grid_colors != nullptr) out_c = torch::empty({0, 3}, vert_options);
             return std::make_tuple(
                 torch::empty({0, 3}, vert_options),
                 torch::empty({0, 3}, tri_options),
-                out_n
+                out_n,
+                out_c,
+                std::nullopt
             );
         }
 
@@ -493,7 +516,15 @@ namespace mc
             out_normals_opt = out_normals;
         }
 
-        interpolate_vertices(out_num_vertices, unique_edges, grid_vertices, voxel_values, grid_normals, iso, p_out_vertices, p_out_normals);
+        std::optional<torch::Tensor> out_colors_opt = std::nullopt;
+        float3* __restrict__ p_out_colors = nullptr;
+        if (grid_colors != nullptr) {
+            torch::Tensor out_colors = torch::empty({out_num_vertices, 3}, vert_options);
+            p_out_colors = (float3*)out_colors.data_ptr<float>();
+            out_colors_opt = out_colors;
+        }
+
+        interpolate_vertices(out_num_vertices, unique_edges, grid_vertices, voxel_values, grid_normals, grid_colors, iso, p_out_vertices, p_out_normals, p_out_colors);
 
         uint32_t out_num_triangles;
         uint32_t *__restrict__ voxel_triangle_prefix_sums;
@@ -505,6 +536,13 @@ namespace mc
 
         assemble_triangles(num_active_voxels, used_voxel_codes, voxel_edge_to_vert_idx, voxel_triangle_prefix_sums, p_out_triangles);
 
+        std::optional<torch::Tensor> out_unique_edges_opt = std::nullopt;
+        if (return_unique_edges) {
+            torch::Tensor out_unique_edges = torch::empty({out_num_vertices, 2}, torch::dtype(torch::kInt32).device(vert_options.device()));
+            CHECK_CUDA_INTERNAL(cudaMemcpy(out_unique_edges.data_ptr(), unique_edges, out_num_vertices * sizeof(Edge), cudaMemcpyDeviceToDevice));
+            out_unique_edges_opt = out_unique_edges;
+        }
+
         // Cleanup
         CHECK_CUDA_INTERNAL(cudaFree(voxel_codes));
         CHECK_CUDA_INTERNAL(cudaFree(used_voxel_index));
@@ -514,6 +552,112 @@ namespace mc
         CHECK_CUDA_INTERNAL(cudaFree(voxel_edge_to_vert_idx));
         CHECK_CUDA_INTERNAL(cudaFree(voxel_triangle_prefix_sums));
 
-        return std::make_tuple(out_vertices, out_triangles, out_normals_opt);
+        return std::make_tuple(out_vertices, out_triangles, out_normals_opt, out_colors_opt, out_unique_edges_opt);
+    }
+
+    __global__ void backward_dmc_kernel(
+        const uint32_t n_verts,
+        const Edge *unique_edges,
+        const float *grid_values,
+        const float3 *grid_coords,
+        const float3 *grid_colors,
+        const float3 *adj_verts,
+        const float3 *adj_colors,
+        const float iso,
+        float *adj_values,
+        float3 *adj_grid_colors,
+        bool with_colors)
+    {
+        uint32_t v_idx = blockIdx.x * blockDim.x + threadIdx.x;
+        if (v_idx >= n_verts)
+            return;
+
+        Edge edge_sig = unique_edges[v_idx];
+        uint32_t v0_idx = edge_sig.v0;
+        uint32_t v1_idx = edge_sig.v1;
+
+        float v0_val = grid_values[v0_idx];
+        float v1_val = grid_values[v1_idx];
+        float3 p0 = grid_coords[v0_idx];
+        float3 p1 = grid_coords[v1_idx];
+
+        float3 grad_p_out = adj_verts[v_idx];
+
+        float3 c0, c1, grad_c_out;
+        if (with_colors)
+        {
+            c0 = grid_colors[v0_idx];
+            c1 = grid_colors[v1_idx];
+            grad_c_out = adj_colors[v_idx];
+        }
+
+        float diff = v1_val - v0_val;
+
+        if (with_colors)
+        {
+            float t = 0.5f;
+            if (fabsf(iso - v0_val) < EPS)
+                t = 0.0f;
+            else if (fabsf(iso - v1_val) < EPS)
+                t = 1.0f;
+            else if (fabsf(diff) >= EPS)
+            {
+                t = fmaxf(0.0f, fminf(1.0f, (iso - v0_val) / diff));
+            }
+
+            float t0 = 1.0f - t;
+            float t1 = t;
+            
+            atomicAdd(&(adj_grid_colors[v0_idx].x), grad_c_out.x * t0);
+            atomicAdd(&(adj_grid_colors[v0_idx].y), grad_c_out.y * t0);
+            atomicAdd(&(adj_grid_colors[v0_idx].z), grad_c_out.z * t0);
+
+            atomicAdd(&(adj_grid_colors[v1_idx].x), grad_c_out.x * t1);
+            atomicAdd(&(adj_grid_colors[v1_idx].y), grad_c_out.y * t1);
+            atomicAdd(&(adj_grid_colors[v1_idx].z), grad_c_out.z * t1);
+        }
+
+        if (diff * diff < 1e-14f)
+            return;
+            
+        float dot_prod = (p1.x - p0.x) * grad_p_out.x + (p1.y - p0.y) * grad_p_out.y + (p1.z - p0.z) * grad_p_out.z;
+        float common = dot_prod / (diff * diff);
+        float grad_v0 = common * (iso - v1_val);
+        float grad_v1 = common * (v0_val - iso);
+
+        atomicAdd(&adj_values[v0_idx], grad_v0);
+        atomicAdd(&adj_values[v1_idx], grad_v1);
+    }
+
+    void backward(
+        const uint32_t n_verts,
+        const Edge *unique_edges,
+        const float3 *grid_vertices,
+        const float3 *grid_colors,
+        const float *values,
+        const float3 *adj_verts,
+        const float3 *adj_colors,
+        float *adj_values,
+        float3 *adj_grid_colors,
+        const float iso)
+    {
+        if (n_verts == 0) return;
+
+        bool with_colors = (grid_colors != nullptr && adj_colors != nullptr && adj_grid_colors != nullptr);
+        int block_size = NTHREADS;
+        int grid_size = (n_verts + block_size - 1) / block_size;
+        
+        backward_dmc_kernel<<<grid_size, block_size>>>(
+            n_verts,
+            unique_edges,
+            values,
+            grid_vertices,
+            grid_colors,
+            adj_verts,
+            adj_colors,
+            iso,
+            adj_values,
+            adj_grid_colors,
+            with_colors);
     }
 }

conquer3d-0.3.0/conquer3d/csrc/ops/mc.h

@@ -0,0 +1,39 @@
+#ifndef MC_H
+#define MC_H
+
+#include <torch/extension.h>
+#include "mc_data.h"
+#include "../check.h"
+#include "../constants.h"
+#include "../primitive/edge.h"
+
+#include <optional>
+
+namespace mc{
+    std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Tensor>, std::optional<torch::Tensor>> marching_cubes(
+        const uint32_t num_voxels,
+        const float3* __restrict__ grid_vertices,
+        const uint32_t* __restrict__ voxels,
+        const float* __restrict__ voxel_values,
+        const float3* __restrict__ grid_normals,
+        const float3* __restrict__ grid_colors,
+        const float iso,
+        torch::TensorOptions vert_options,
+        torch::TensorOptions tri_options,
+        bool return_unique_edges = false
+    );
+    void backward(
+        const uint32_t n_verts,
+        const Edge *unique_edges,
+        const float3 *grid_vertices,
+        const float3 *grid_colors,
+        const float *values,
+        const float3 *adj_verts,
+        const float3 *adj_colors,
+        float *adj_values,
+        float3 *adj_grid_colors,
+        const float iso
+    );
+}
+
+#endif // MC_H

conquer3d-0.3.0/conquer3d/data_structure/__init__.py

@@ -0,0 +1,15 @@
+from .._C import (
+    KDTree,
+    BVH,
+    GSBVH,
+    PGSBVH,
+    MeshBVH,
+    TriangleMesh
+)
+from .grid import create_voxel_grid, compute_grid_normal, compute_active_voxels, voxel2sparse, sparse2voxel
+
+spatial_data_structures = ['KDTree', 'BVH', 'GSBVH', 'PGSBVH', 'MeshBVH']
+mesh_data_structures = ['TriangleMesh']
+grid_data_structures = ['create_voxel_grid', 'compute_grid_normal', 'compute_active_voxels', 'voxel2sparse', 'sparse2voxel']
+
+__all__ = spatial_data_structures + mesh_data_structures + grid_data_structures

conquer3d-0.3.0/conquer3d/data_structure/grid.py

@@ -0,0 +1,114 @@
+import torch
+from typing import Tuple, List, Union
+from .. import _C
+
+def create_voxel_grid(
+    grid_min: Union[List[float], Tuple[float, float, float]],
+    grid_max: Union[List[float], Tuple[float, float, float]],
+    res: Union[List[int], Tuple[int, int, int]],
+    device: str = "cuda"
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Creates a structured 3D voxel grid efficiently.
+    
+    Args:
+        grid_min (List[float] | Tuple[float, float, float]): The minimum (x, y, z) bounding box coordinates.
+        grid_max (List[float] | Tuple[float, float, float]): The maximum (x, y, z) bounding box coordinates.
+        res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
+        device (str, optional): Target device for the tensors (e.g., "cuda" or "cpu"). Defaults to "cuda".
+        
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            - grid_vertices (torch.Tensor): A float32 tensor of shape (N, 3) containing all grid coordinates.
+            - voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices 
+                                     for each voxel, mapped identically to the Marching Cubes convention.
+            - idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) 3D coordinate indices for each vertex.
+    """
+    return _C.create_voxel_grid(list(grid_min), list(grid_max), list(res), device)
+
+def compute_grid_normal(
+    sdf: torch.Tensor,
+    grid_vertices: torch.Tensor,
+    idx_grids: torch.Tensor,
+    res: Union[List[int], Tuple[int, int, int]]
+) -> torch.Tensor:
+    """
+    Computes smooth normals (gradients) for a voxel grid using central differences.
+    
+    Args:
+        sdf (torch.Tensor): A float32 tensor of shape (N,) containing the SDF values.
+        grid_vertices (torch.Tensor): A float32 tensor of shape (N, 3) containing the grid coordinates.
+        idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) indices, returned by create_voxel_grid.
+        res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
+        
+    Returns:
+        torch.Tensor: A float32 tensor of shape (N, 3) containing the normalized gradient vectors.
+    """
+    return _C.compute_grid_normal(sdf.contiguous(), grid_vertices.contiguous(), idx_grids.contiguous(), list(res))
+
+def compute_active_voxels(
+    voxels: torch.Tensor,
+    sdf: torch.Tensor,
+    iso: float = 0.0
+) -> torch.Tensor:
+    """
+    Computes the indices of active voxels that intersect the isosurface.
+    
+    Args:
+        voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices for each voxel.
+        sdf (torch.Tensor): A float32 tensor of shape (N,) containing the SDF values.
+        iso (float): The isosurface value. Defaults to 0.0.
+        
+    Returns:
+        torch.Tensor: An int64 tensor containing the indices of the active voxels.
+    """
+    return _C.compute_active_voxels(voxels.contiguous(), sdf.contiguous(), iso)
+
+def voxel2sparse(
+    active_indices: torch.Tensor,
+    voxels: torch.Tensor,
+    idx_grids: torch.Tensor,
+    batch_idx: int = 0
+) -> torch.Tensor:
+    """
+    Converts 1D active voxel indices into the [N, 4] coordinates tensor 
+    expected by sparse neural networks (like spconv/torchsparse).
+    
+    Args:
+        active_indices (torch.Tensor): 1D tensor of active voxel indices.
+        voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices for each voxel.
+        idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) 3D coordinate indices for each vertex.
+        batch_idx (int, optional): The batch index to prepend. Defaults to 0.
+        
+    Returns:
+        torch.Tensor: An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
+    """
+    active_corners = voxels[active_indices]
+    base_vertex_indices = active_corners[:, 0].to(torch.int64)
+    voxel_coords = idx_grids[base_vertex_indices]
+    
+    batch_col = torch.full((voxel_coords.shape[0], 1), batch_idx, device=voxel_coords.device, dtype=torch.int64)
+    sparse_coords = torch.cat([batch_col, voxel_coords], dim=1).to(torch.int32)
+    return sparse_coords
+
+def sparse2voxel(
+    sparse_coords: torch.Tensor,
+    res: Union[List[int], Tuple[int, int, int]]
+) -> torch.Tensor:
+    """
+    Converts a [N, 4] sparse coordinates tensor back to 1D active voxel indices.
+    
+    Args:
+        sparse_coords (torch.Tensor): An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
+        res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
+        
+    Returns:
+        torch.Tensor: An int64 tensor containing the indices of the active voxels.
+    """
+    x = sparse_coords[:, 1].to(torch.int64)
+    y = sparse_coords[:, 2].to(torch.int64)
+    z = sparse_coords[:, 3].to(torch.int64)
+    
+    rx, ry, rz = res
+    active_indices = x * (ry - 1) * (rz - 1) + y * (rz - 1) + z
+    return active_indices

{conquer3d-0.2.9 → conquer3d-0.3.0}/conquer3d/ops/__init__.py

@@ -1,11 +1,10 @@
 from .marching_cubes import marching_cubes
-from .grid import create_voxel_grid, compute_grid_normal
+from .diff_marching_cubes import diff_marching_cubes
 from .distance import one_sided_chamfer_distance, chamfer_distance
 
 __all__ = [
     "marching_cubes",
-    "create_voxel_grid",
-    "compute_grid_normal",
+    "diff_marching_cubes",
     "one_sided_chamfer_distance",
     "chamfer_distance"
 ]

conquer3d-0.3.0/conquer3d/ops/diff_marching_cubes.py

@@ -0,0 +1,135 @@
+import torch
+from typing import Tuple, Optional
+
+# Explicitly import the compiled CMake target
+from .._C import marching_cubes as marching_cubes_func
+from .._C import marching_cubes_backward as marching_cubes_backward_func
+
+class DiffMarchingCubes(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        grid_vertices: torch.Tensor,
+        voxels: torch.Tensor,
+        voxel_values: torch.Tensor,
+        grid_normals: Optional[torch.Tensor] = None,
+        grid_colors: Optional[torch.Tensor] = None,
+        iso: float = 0.0
+    ):
+        # We need return_unique_edges = True for the backward pass
+        out_vertices, out_triangles, out_normals, out_colors, unique_edges = marching_cubes_func(
+            grid_vertices,
+            voxels,
+            voxel_values,
+            grid_normals,
+            grid_colors,
+            iso,
+            True # return_unique_edges
+        )
+        
+        ctx.iso = iso
+        ctx.has_colors = grid_colors is not None
+        
+        # Save tensors needed for the backward pass
+        if ctx.has_colors:
+            ctx.save_for_backward(unique_edges, grid_vertices, voxel_values, grid_colors)
+        else:
+            ctx.save_for_backward(unique_edges, grid_vertices, voxel_values)
+        
+        return out_vertices, out_triangles, out_normals, out_colors
+
+    @staticmethod
+    def backward(ctx, grad_out_vertices, grad_out_triangles, grad_out_normals, grad_out_colors):
+        if ctx.has_colors:
+            unique_edges, grid_vertices, voxel_values, grid_colors = ctx.saved_tensors
+        else:
+            unique_edges, grid_vertices, voxel_values = ctx.saved_tensors
+            grid_colors = None
+            
+        iso = ctx.iso
+        
+        # Initialize gradients for inputs
+        grad_voxel_values = torch.zeros_like(voxel_values)
+        
+        grad_grid_colors = None
+        if ctx.has_colors:
+            grad_grid_colors = torch.zeros_like(grid_colors)
+            
+        # Call the backward C++ binding if we have valid output vertices
+        if unique_edges is not None and unique_edges.shape[0] > 0 and grad_out_vertices is not None:
+            grad_out_vertices = grad_out_vertices.contiguous()
+            if grad_out_colors is not None:
+                grad_out_colors = grad_out_colors.contiguous()
+                
+            marching_cubes_backward_func(
+                unique_edges,
+                grid_vertices,
+                grid_colors,
+                voxel_values,
+                grad_out_vertices,
+                grad_out_colors,
+                grad_voxel_values,
+                grad_grid_colors,
+                iso
+            )
+            
+        # Return gradients for all inputs in the same order as forward
+        # grid_vertices, voxels, voxel_values, grid_normals, grid_colors, iso
+        return None, None, grad_voxel_values, None, grad_grid_colors, None
+
+
+def diff_marching_cubes(
+    grid_vertices: torch.Tensor,
+    voxels: torch.Tensor,
+    voxel_values: torch.Tensor,
+    grid_normals: Optional[torch.Tensor] = None,
+    grid_colors: Optional[torch.Tensor] = None,
+    iso: float = 0.0
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+    """
+    Executes the Differentiable Marching Cubes algorithm to extract an isosurface from a voxel grid.
+
+    Args:
+        grid_vertices (torch.Tensor): (V, 3) tensor of global vertex positions.
+        voxels (torch.Tensor): (N, 8) tensor of voxel corner indices mapping to `grid_vertices`.
+        voxel_values (torch.Tensor): (V,) tensor of SDF/scalar values at each vertex. Requires grad.
+        grid_normals (torch.Tensor, optional): (V, 3) optional tensor of SDF normals at each vertex. Defaults to None.
+        grid_colors (torch.Tensor, optional): (V, 3) optional tensor of RGB colors at each vertex. Defaults to None. Requires grad.
+        iso (float, optional): The isosurface extraction threshold. Defaults to 0.0.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+            - vertices (torch.Tensor): (M, 3) tensor of extracted mesh vertices. Supports gradients.
+            - triangles (torch.Tensor): (T, 3) tensor of extracted mesh triangle indices.
+            - normals (torch.Tensor, optional): (M, 3) tensor of extracted mesh vertex normals, if `grid_normals` was provided.
+            - colors (torch.Tensor, optional): (M, 3) tensor of extracted mesh vertex colors, if `grid_colors` was provided. Supports gradients.
+    """
+    if not all(t.is_cuda for t in [grid_vertices, voxels, voxel_values]):
+        raise ValueError("All input tensors must be CUDA tensors.")
+        
+    grid_vertices_c = grid_vertices.contiguous().to(torch.float32)
+    voxels_c = voxels.contiguous().to(torch.int32)
+    voxel_values_c = voxel_values.contiguous().to(torch.float32)
+    
+    grid_normals_c = None
+    if grid_normals is not None:
+        if not grid_normals.is_cuda:
+            raise ValueError("grid_normals must be a CUDA tensor.")
+        grid_normals_c = grid_normals.contiguous().to(torch.float32)
+
+    grid_colors_c = None
+    if grid_colors is not None:
+        if not grid_colors.is_cuda:
+            raise ValueError("grid_colors must be a CUDA tensor.")
+        grid_colors_c = grid_colors.contiguous().to(torch.float32)
+
+    iso_c = float(iso)
+
+    return DiffMarchingCubes.apply(
+        grid_vertices_c,
+        voxels_c,
+        voxel_values_c,
+        grid_normals_c,
+        grid_colors_c,
+        iso_c
+    )

{conquer3d-0.2.9 → conquer3d-0.3.0}/conquer3d/ops/marching_cubes.py

@@ -9,8 +9,9 @@ def marching_cubes(
     voxels: torch.Tensor,
     voxel_values: torch.Tensor,
     grid_normals: Optional[torch.Tensor] = None,
+    grid_colors: Optional[torch.Tensor] = None,
     iso: float = 0.0
-) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
     """
     Executes the Marching Cubes algorithm to extract an isosurface from a voxel grid.
 
@@ -19,13 +20,15 @@ def marching_cubes(
         voxels (torch.Tensor): (N, 8) tensor of voxel corner indices mapping to `grid_vertices`.
         voxel_values (torch.Tensor): (V,) tensor of SDF/scalar values at each vertex.
         grid_normals (torch.Tensor, optional): (V, 3) optional tensor of SDF normals at each vertex. Defaults to None.
+        grid_colors (torch.Tensor, optional): (V, 3) optional tensor of RGB colors at each vertex. Defaults to None.
         iso (float, optional): The isosurface extraction threshold. Defaults to 0.0.
 
     Returns:
-        Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
             - vertices (torch.Tensor): (M, 3) tensor of extracted mesh vertices.
             - triangles (torch.Tensor): (T, 3) tensor of extracted mesh triangle indices.
             - normals (torch.Tensor, optional): (M, 3) tensor of extracted mesh vertex normals, if `grid_normals` was provided.
+            - colors (torch.Tensor, optional): (M, 3) tensor of extracted mesh vertex colors, if `grid_colors` was provided.
     """
     if not all(t.is_cuda for t in [grid_vertices, voxels, voxel_values]):
         raise ValueError("All input tensors must be CUDA tensors.")
@@ -40,14 +43,22 @@ def marching_cubes(
             raise ValueError("grid_normals must be a CUDA tensor.")
         grid_normals_c = grid_normals.contiguous().to(torch.float32)
 
+    grid_colors_c = None
+    if grid_colors is not None:
+        if not grid_colors.is_cuda:
+            raise ValueError("grid_colors must be a CUDA tensor.")
+        grid_colors_c = grid_colors.contiguous().to(torch.float32)
+
     iso_c = float(iso)
 
-    vertices, triangles, out_normals = marching_cubes_func(
+    vertices, triangles, out_normals, out_colors, _ = marching_cubes_func(
         grid_vertices_c,
         voxels_c,
         voxel_values_c,
         grid_normals_c,
-        iso_c
+        grid_colors_c,
+        iso_c,
+        False # return_unique_edges
     )
 
-    return vertices, triangles, out_normals
+    return vertices, triangles, out_normals, out_colors

{conquer3d-0.2.9 → conquer3d-0.3.0/conquer3d.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.9
+Version: 0.3.0
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT
@@ -51,8 +51,13 @@ conda activate geocutool
 
 conda install nvidia::cuda-toolkit==12.8.2 -y
 
+pip install setuptools wheel ninja
 pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
 
+pip install git+https://github.com/mit-han-lab/torchsparse.git
+
+pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.8.0_cu128.html
+
 pip install pybind11-stubgen
 pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
 

{conquer3d-0.2.9 → conquer3d-0.3.0}/conquer3d.egg-info/SOURCES.txt

@@ -67,9 +67,10 @@ conquer3d/csrc/primitive/pgs_math.cuh
 conquer3d/csrc/primitive/ray.h
 conquer3d/csrc/primitive/triangles.h
 conquer3d/data_structure/__init__.py
+conquer3d/data_structure/grid.py
 conquer3d/ops/__init__.py
+conquer3d/ops/diff_marching_cubes.py
 conquer3d/ops/distance.py
-conquer3d/ops/grid.py
 conquer3d/ops/marching_cubes.py
 conquer3d/primitive/__init__.py
 conquer3d/primitive/gs.py

{conquer3d-0.2.9 → conquer3d-0.3.0}/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "conquer3d"
-version = "0.2.9"
+version = "0.3.0"
 description = "Geometric Cuda Tool Box"
 readme = "README.md"
 license = "MIT"

conquer3d-0.2.9/conquer3d/csrc/binds/ops/mc.cpp

@@ -1,47 +0,0 @@
-#include <torch/extension.h>
-#include <pybind11/pybind11.h>
-#include "../../ops/mc.h"
-#include "../../check.h"
-
-namespace py = pybind11;
-
-std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> marching_cubes_wrapper(
-    torch::Tensor grid_vertices,
-    torch::Tensor voxels,
-    torch::Tensor voxel_values,
-    std::optional<torch::Tensor> grid_normals,
-    float iso
-) {
-    CHECK_INPUT(grid_vertices);
-    CHECK_INPUT(voxels);
-    CHECK_INPUT(voxel_values);
-
-    uint32_t num_voxels = voxels.size(0);
-
-    const float3* __restrict__ p_grid_vertices = (float3*)grid_vertices.data_ptr<float>();
-    const uint32_t* __restrict__ p_voxels = (uint32_t*)voxels.data_ptr<int32_t>();
-    const float* __restrict__ p_voxel_values = voxel_values.data_ptr<float>();
-
-    const float3* __restrict__ p_grid_normals = nullptr;
-    if (grid_normals.has_value()) {
-        CHECK_INPUT(grid_normals.value());
-        p_grid_normals = (float3*)grid_normals.value().data_ptr<float>();
-    }
-
-    return mc::marching_cubes(
-        num_voxels,
-        p_grid_vertices,
-        p_voxels,
-        p_voxel_values,
-        p_grid_normals,
-        iso,
-        grid_vertices.options(),
-        voxels.options()
-    );
-}
-
-void bind_ops_mc(py::module_& m) {
-    m.def("marching_cubes", &marching_cubes_wrapper, "Marching Cubes",
-          py::arg("grid_vertices"), py::arg("voxels"), py::arg("voxel_values"),
-          py::arg("grid_normals") = std::nullopt, py::arg("iso") = 0.0f);
-}

conquer3d-0.2.9/conquer3d/csrc/ops/mc.h

@@ -1,25 +0,0 @@
-#ifndef MC_H
-#define MC_H
-
-#include <torch/extension.h>
-#include "mc_data.h"
-#include "../check.h"
-#include "../constants.h"
-#include "../primitive/edge.h"
-
-#include <optional>
-
-namespace mc{
-    std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> marching_cubes(
-        const uint32_t num_voxels,
-        const float3* __restrict__ grid_vertices,
-        const uint32_t* __restrict__ voxels,
-        const float* __restrict__ voxel_values,
-        const float3* __restrict__ grid_normals,
-        const float iso,
-        torch::TensorOptions vert_options,
-        torch::TensorOptions tri_options
-    );
-}
-
-#endif // MC_H

conquer3d-0.2.9/conquer3d/data_structure/__init__.py

@@ -1,13 +0,0 @@
-from .._C import (
-    KDTree,
-    BVH,
-    GSBVH,
-    PGSBVH,
-    MeshBVH,
-    TriangleMesh
-)
-
-spatial_data_structures = ['KDTree', 'BVH', 'GSBVH', 'PGSBVH', 'MeshBVH']
-mesh_data_structures = ['TriangleMesh']
-
-__all__ = spatial_data_structures + mesh_data_structures

conquer3d-0.2.9/conquer3d/ops/grid.py

@@ -1,47 +0,0 @@
-import torch
-from typing import Tuple, List, Union
-from .. import _C
-
-def create_voxel_grid(
-    grid_min: Union[List[float], Tuple[float, float, float]],
-    grid_max: Union[List[float], Tuple[float, float, float]],
-    res: Union[List[int], Tuple[int, int, int]],
-    device: str = "cuda"
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Creates a structured 3D voxel grid efficiently.
-    
-    Args:
-        grid_min (List[float] | Tuple[float, float, float]): The minimum (x, y, z) bounding box coordinates.
-        grid_max (List[float] | Tuple[float, float, float]): The maximum (x, y, z) bounding box coordinates.
-        res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
-        device (str, optional): Target device for the tensors (e.g., "cuda" or "cpu"). Defaults to "cuda".
-        
-    Returns:
-        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            - grid_vertices (torch.Tensor): A float32 tensor of shape (N, 3) containing all grid coordinates.
-            - voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices 
-                                     for each voxel, mapped identically to the Marching Cubes convention.
-            - idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) 3D coordinate indices for each vertex.
-    """
-    return _C.create_voxel_grid(list(grid_min), list(grid_max), list(res), device)
-
-def compute_grid_normal(
-    sdf: torch.Tensor,
-    grid_vertices: torch.Tensor,
-    idx_grids: torch.Tensor,
-    res: Union[List[int], Tuple[int, int, int]]
-) -> torch.Tensor:
-    """
-    Computes smooth normals (gradients) for a voxel grid using central differences.
-    
-    Args:
-        sdf (torch.Tensor): A float32 tensor of shape (N,) containing the SDF values.
-        grid_vertices (torch.Tensor): A float32 tensor of shape (N, 3) containing the grid coordinates.
-        idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) indices, returned by create_voxel_grid.
-        res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
-        
-    Returns:
-        torch.Tensor: A float32 tensor of shape (N, 3) containing the normalized gradient vectors.
-    """
-    return _C.compute_grid_normal(sdf.contiguous(), grid_vertices.contiguous(), idx_grids.contiguous(), list(res))