conquer3d 0.2.8__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.8
+Version: 0.3.0
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT
@@ -51,12 +51,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
+pip install plotly open3d jupyter trimesh point-cloud-utils
 ```
 
 # Development

{conquer3d-0.2.8 → conquer3d-0.3.0}/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
+pip install plotly open3d jupyter trimesh point-cloud-utils
 ```
 
 # Development

conquer3d-0.3.0/conquer3d/csrc/binds/data_structure/grid.cpp

@@ -0,0 +1,150 @@
+#include <torch/extension.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <tuple>
+#include <string>
+#include <vector>
+
+namespace py = pybind11;
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid(
+    std::vector<float> grid_min,
+    std::vector<float> grid_max,
+    std::vector<int64_t> res,
+    std::string device_str
+) {
+    TORCH_CHECK(grid_min.size() == 3, "grid_min must have 3 elements.");
+    TORCH_CHECK(grid_max.size() == 3, "grid_max must have 3 elements.");
+    TORCH_CHECK(res.size() == 3, "res must have 3 elements.");
+
+    c10::Device device(device_str);
+    auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+    int64_t rx = res[0];
+    int64_t ry = res[1];
+    int64_t rz = res[2];
+
+    auto x = torch::linspace(grid_min[0], grid_max[0], rx, options);
+    auto y = torch::linspace(grid_min[1], grid_max[1], ry, options);
+    auto z = torch::linspace(grid_min[2], grid_max[2], rz, options);
+
+    auto grids = torch::meshgrid({x, y, z}, "ij");
+    auto grid_x = grids[0].flatten();
+    auto grid_y = grids[1].flatten();
+    auto grid_z = grids[2].flatten();
+    
+    auto grid_vertices = torch::stack({grid_x, grid_y, grid_z}, 1).contiguous();
+
+    auto options_int = options.dtype(torch::kInt64);
+    auto i = torch::arange(rx - 1, options_int);
+    auto j = torch::arange(ry - 1, options_int);
+    auto k = torch::arange(rz - 1, options_int);
+
+    auto idx_grids = torch::meshgrid({i, j, k}, "ij");
+    auto vi = idx_grids[0];
+    auto vj = idx_grids[1];
+    auto vk = idx_grids[2];
+
+    // Voxel connectivity indexing mapping matching mc.cu
+    auto v0 = vi * ry * rz + vj * rz + vk;
+    auto v1 = (vi + 1) * ry * rz + vj * rz + vk;
+    auto v2 = (vi + 1) * ry * rz + (vj + 1) * rz + vk;
+    auto v3 = vi * ry * rz + (vj + 1) * rz + vk;
+    auto v4 = vi * ry * rz + vj * rz + (vk + 1);
+    auto v5 = (vi + 1) * ry * rz + vj * rz + (vk + 1);
+    auto v6 = (vi + 1) * ry * rz + (vj + 1) * rz + (vk + 1);
+    auto v7 = vi * ry * rz + (vj + 1) * rz + (vk + 1);
+
+    auto voxels = torch::stack({v0, v1, v2, v3, v4, v5, v6, v7}, -1).view({-1, 8}).contiguous().to(torch::kInt32);
+
+    auto i_v = torch::arange(rx, options_int);
+    auto j_v = torch::arange(ry, options_int);
+    auto k_v = torch::arange(rz, options_int);
+    auto idx_grids_v = torch::meshgrid({i_v, j_v, k_v}, "ij");
+    auto out_idx_grids = torch::stack({idx_grids_v[0].flatten(), idx_grids_v[1].flatten(), idx_grids_v[2].flatten()}, -1).contiguous();
+
+    return std::make_tuple(grid_vertices, voxels, out_idx_grids);
+}
+
+torch::Tensor compute_grid_normal(torch::Tensor sdf, torch::Tensor grid_vertices, torch::Tensor idx_grids, std::vector<int64_t> res) {
+    TORCH_CHECK(res.size() == 3, "res must have 3 elements.");
+    
+    auto i = idx_grids.select(1, 0);
+    auto j = idx_grids.select(1, 1);
+    auto k = idx_grids.select(1, 2);
+    
+    auto rx = res[0];
+    auto ry = res[1];
+    auto rz = res[2];
+    
+    auto i_prev = torch::clamp(i - 1, 0, rx - 1);
+    auto i_next = torch::clamp(i + 1, 0, rx - 1);
+    
+    auto j_prev = torch::clamp(j - 1, 0, ry - 1);
+    auto j_next = torch::clamp(j + 1, 0, ry - 1);
+    
+    auto k_prev = torch::clamp(k - 1, 0, rz - 1);
+    auto k_next = torch::clamp(k + 1, 0, rz - 1);
+    
+    auto idx_x_prev = i_prev * ry * rz + j * rz + k;
+    auto idx_x_next = i_next * ry * rz + j * rz + k;
+    
+    auto idx_y_prev = i * ry * rz + j_prev * rz + k;
+    auto idx_y_next = i * ry * rz + j_next * rz + k;
+    
+    auto idx_z_prev = i * ry * rz + j * rz + k_prev;
+    auto idx_z_next = i * ry * rz + j * rz + k_next;
+    
+    auto sdf_flat = sdf.flatten();
+    
+    auto diff_x = sdf_flat.index_select(0, idx_x_next) - sdf_flat.index_select(0, idx_x_prev);
+    auto diff_y = sdf_flat.index_select(0, idx_y_next) - sdf_flat.index_select(0, idx_y_prev);
+    auto diff_z = sdf_flat.index_select(0, idx_z_next) - sdf_flat.index_select(0, idx_z_prev);
+
+    auto vx = grid_vertices.select(1, 0);
+    auto vy = grid_vertices.select(1, 1);
+    auto vz = grid_vertices.select(1, 2);
+
+    auto dx = vx.index_select(0, idx_x_next) - vx.index_select(0, idx_x_prev);
+    auto dy = vy.index_select(0, idx_y_next) - vy.index_select(0, idx_y_prev);
+    auto dz = vz.index_select(0, idx_z_next) - vz.index_select(0, idx_z_prev);
+
+    dx = torch::where(dx == 0, torch::tensor(1e-5f, dx.options()), dx);
+    dy = torch::where(dy == 0, torch::tensor(1e-5f, dy.options()), dy);
+    dz = torch::where(dz == 0, torch::tensor(1e-5f, dz.options()), dz);
+
+    auto grad_x = diff_x / dx;
+    auto grad_y = diff_y / dy;
+    auto grad_z = diff_z / dz;
+    
+    auto normals = torch::stack({grad_x, grad_y, grad_z}, 1);
+    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/chamfer.cpp

@@ -0,0 +1,42 @@
+#include <torch/extension.h>
+#include <pybind11/pybind11.h>
+#include "../../ops/chamfer.h"
+#include "../../check.h"
+
+namespace py = pybind11;
+
+std::tuple<torch::Tensor, torch::Tensor> one_sided_chamfer_distance_wrapper(
+    torch::Tensor query_points,
+    torch::Tensor reference_points
+) {
+    CHECK_INPUT(query_points);
+    CHECK_INPUT(reference_points);
+
+    uint32_t num_query_points = query_points.size(0);
+    uint32_t num_reference_points = reference_points.size(0);
+
+    const float3* __restrict__ p_query_points = (float3*)query_points.data_ptr<float>();
+    const float3* __restrict__ p_reference_points = (float3*)reference_points.data_ptr<float>();
+
+    torch::Tensor distances = torch::empty({num_query_points}, query_points.options());
+    torch::Tensor indices = torch::empty({num_query_points}, query_points.options().dtype(torch::kInt64));
+
+    float* __restrict__ p_distances = distances.data_ptr<float>();
+    int64_t* __restrict__ p_indices = indices.data_ptr<int64_t>();
+
+    one_sided_chamfer_distance(
+        num_query_points,
+        p_query_points,
+        num_reference_points,
+        p_reference_points,
+        p_distances,
+        p_indices
+    );
+
+    return std::make_tuple(distances, indices);
+}
+
+void bind_ops_chamfer(py::module_& m) {
+    m.def("one_sided_chamfer_distance", &one_sided_chamfer_distance_wrapper, "One Sided Chamfer Distance",
+          py::arg("query_points"), py::arg("reference_points"));
+}

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.3.0/conquer3d/csrc/ops/chamfer.cu

@@ -0,0 +1,73 @@
+#include "../data_structure/kdtree.h"
+
+#include <cstdint>
+#include <thrust/device_vector.h>
+#include <thrust/sequence.h>
+
+__global__ void one_sided_chamfer_distance_kernel(
+    const uint32_t num_query_points,
+    const float3 *__restrict__ query_points,
+    const uint32_t num_reference_points,
+    const float3 *__restrict__ tree_points,
+    const int64_t *__restrict__ tree_inds,
+    float *__restrict__ distances,
+    int64_t *__restrict__ indices)
+{
+    uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= num_query_points)
+        return;
+
+    float3 query_point = query_points[idx];
+
+    float best_dists[MAX_K];
+    int64_t best_inds[MAX_K];
+
+#pragma unroll
+    for (int i = 0; i < MAX_K; i++)
+    {
+        best_dists[i] = FLT_MAX;
+        best_inds[i] = -1;
+    }
+
+    kdtree::query_kdtree_loop(
+        query_point,
+        num_reference_points,
+        tree_points,
+        tree_inds,
+        1,
+        best_dists,
+        best_inds);
+    
+    distances[idx] = best_dists[0];
+    indices[idx] = best_inds[0];
+}
+
+void one_sided_chamfer_distance(
+    const uint32_t num_query_points,
+    const float3 *__restrict__ query_points,
+    const uint32_t num_reference_points,
+    const float3 *__restrict__ reference_points,
+    float *__restrict__ distances,
+    int64_t *__restrict__ indices)
+{
+    thrust::device_vector<float3> cloned_ref_points(reference_points, reference_points + num_reference_points);
+    thrust::device_vector<int64_t> reference_indices(num_reference_points);
+    thrust::sequence(reference_indices.begin(), reference_indices.end());
+
+    kdtree::build(
+        num_reference_points,
+        thrust::raw_pointer_cast(cloned_ref_points.data()),
+        thrust::raw_pointer_cast(reference_indices.data()));
+
+    uint32_t threads = NTHREADS;
+    uint32_t blocks = (num_query_points + threads - 1) / threads;
+
+    one_sided_chamfer_distance_kernel<<<blocks, threads>>>(
+        num_query_points,
+        query_points,
+        num_reference_points,
+        thrust::raw_pointer_cast(cloned_ref_points.data()),
+        thrust::raw_pointer_cast(reference_indices.data()),
+        distances,
+        indices);
+}

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

@@ -0,0 +1,18 @@
+#ifndef CHAMFER_H
+#define CHAMFER_H
+
+#include "../constants.h"
+#include "../maths/maths.h"
+
+#include <cuda_runtime.h>
+
+void one_sided_chamfer_distance(
+    const uint32_t num_query_points,
+    const float3* __restrict__ query_points,
+    const uint32_t num_reference_points,
+    const float3* __restrict__ reference_points,
+    float* __restrict__ distances,
+    int64_t* __restrict__ indices
+);
+
+#endif // CHAMFER_H