conquer3d 0.2.9__tar.gz → 0.3.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.9
+Version: 0.3.1
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT
@@ -13,6 +13,7 @@ Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
+Requires-Dist: gdown
 Provides-Extra: cuda
 Requires-Dist: cuda-toolkit==12.8.2; extra == "cuda"
 Provides-Extra: dev
@@ -51,12 +52,17 @@ 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
 
-pip install plotly open3d jupyter trimesh point-cloud-utils
+pip install plotly open3d jupyter trimesh point-cloud-utils meshlib
 ```
 
 # Development

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

@@ -23,12 +23,17 @@ 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
 
-pip install plotly open3d jupyter trimesh point-cloud-utils
+pip install plotly open3d jupyter trimesh point-cloud-utils meshlib
 ```
 
 # Development

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

@@ -121,9 +121,29 @@ 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.device() == sdf.device(), "voxels and sdf must be on the same device.");
+    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.1/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.1}/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.1/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.2.9 → conquer3d-0.3.1}/conquer3d/csrc/primitive/aabb.h

@@ -167,6 +167,21 @@ namespace aabb
             (target_aabb_max.y) >= query_aabb_max.y &&
             (target_aabb_max.z) >= query_aabb_max.z);
     }
+
+    __device__ __forceinline__ bool test_point_inside_aabb(
+        const float3 &point,
+        const float3 &aabb_min,
+        const float3 &aabb_max
+    )
+    {
+        return (
+            aabb_min.x <= point.x &&
+            aabb_min.y <= point.y &&
+            aabb_min.z <= point.z &&
+            (aabb_max.x) >= point.x &&
+            (aabb_max.y) >= point.y &&
+            (aabb_max.z) >= point.z);
+    }
 }
 
 #endif // AABB_H

conquer3d-0.3.1/conquer3d/data/collate/sparse_tensor.py

@@ -0,0 +1,43 @@
+import torch
+
+def sparse_collate_fn(batch):
+    """
+    Collate function for sparse tensor dataloaders (e.g., MinkowskiEngine, spconv, torchsparse).
+    Takes a batch of samples where each sample is (idx_grids, features, ...).
+    Appends the batch index to the idx_grids to create (batch_idx, x, y, z).
+    
+    Args:
+        batch (list): A list of tuples (idx_grids, features, label)
+        
+    Returns:
+        tuple: (batched_coords, batched_features, batched_labels)
+            batched_coords (torch.Tensor): Shape (Total_N, 4)
+            batched_features (torch.Tensor): Shape (Total_N, ...)
+            batched_labels (torch.Tensor): Shape (Batch_Size,)
+    """
+    batched_coords = []
+    batched_features = []
+    batched_labels = []
+
+    for batch_idx, item in enumerate(batch):
+        idx_grids, features, label = item
+        
+        # Create a batch index column of shape [N, 1]
+        b_col = torch.full((idx_grids.shape[0], 1), batch_idx, dtype=idx_grids.dtype, device=idx_grids.device)
+        
+        # Concatenate to get [N, 4] formatted as (batch_idx, x, y, z)
+        coords_with_batch = torch.cat([b_col, idx_grids], dim=1)
+        
+        batched_coords.append(coords_with_batch)
+        batched_features.append(features)
+        batched_labels.append(label)
+
+    batched_coords = torch.cat(batched_coords, dim=0)
+    batched_features = torch.cat(batched_features, dim=0)
+    
+    if isinstance(batched_labels[0], torch.Tensor):
+        batched_labels = torch.stack(batched_labels, dim=0)
+    else:
+        batched_labels = torch.tensor(batched_labels)
+
+    return batched_coords, batched_features, batched_labels

conquer3d-0.3.1/conquer3d/data/dataset/base_mesh.py

@@ -0,0 +1,17 @@
+from torch.utils.data import Dataset
+from typing import Callable, Optional
+
+class BaseMeshDataset(Dataset):
+    """
+    Base PyTorch Dataset class for Mesh datasets.
+    """
+    def __init__(self, root: str, transform: Optional[Callable] = None):
+        super().__init__()
+        self.root = root
+        self.transform = transform
+
+    def __len__(self) -> int:
+        raise NotImplementedError("Subclasses must implement __len__")
+
+    def __getitem__(self, idx: int):
+        raise NotImplementedError("Subclasses must implement __getitem__")