conquer3d 0.5.2__tar.gz → 0.5.4__tar.gz

{conquer3d-0.5.2/conquer3d.egg-info → conquer3d-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.5.2
+Version: 0.5.4
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT

{conquer3d-0.5.2 → conquer3d-0.5.4}/conquer3d/csrc/binds/data_structure/grid.cpp

@@ -6,14 +6,17 @@
 #include <vector>
 #include "../../data_structure/mesh_bvh.h"
 #include "../../data_structure/triangle_mesh.h"
+#include "../../data_structure/grid.h"
+#include <cuda_runtime.h>
 
 namespace py = pybind11;
 
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid(
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> create_voxel_grid(
     std::vector<float> grid_min,
     std::vector<float> grid_max,
     std::vector<int64_t> res,
-    std::string device_str
+    std::string device_str,
+    bool return_idx_grids
 ) {
     TORCH_CHECK(grid_min.size() == 3, "grid_min must have 3 elements.");
     TORCH_CHECK(grid_max.size() == 3, "grid_max must have 3 elements.");
@@ -59,13 +62,16 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid(
 
     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();
+    std::optional<torch::Tensor> opt_idx_grids = std::nullopt;
+    if (return_idx_grids) {
+        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");
+        opt_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);
+    return std::make_tuple(grid_vertices, voxels, opt_idx_grids);
 }
 
 torch::Tensor compute_grid_normal(torch::Tensor sdf, torch::Tensor grid_vertices, torch::Tensor idx_grids, std::vector<int64_t> res) {
@@ -141,11 +147,12 @@ torch::Tensor compute_active_voxels(torch::Tensor voxels, torch::Tensor sdf, flo
     return active_indices;
 }
 
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid_from_tmesh(
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> create_voxel_grid_from_tmesh(
     std::vector<float> grid_min,
     std::vector<float> grid_max,
     std::vector<int64_t> res,
-    TriangleMesh &tmesh
+    TriangleMesh &tmesh,
+    bool return_unique_vert_ids = true
 ) {
     TORCH_CHECK(grid_min.size() == 3, "grid_min must have 3 elements.");
     TORCH_CHECK(grid_max.size() == 3, "grid_max must have 3 elements.");
@@ -165,7 +172,131 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid_from_t
         return std::make_tuple(
             torch::empty({0, 3}, torch::TensorOptions().dtype(torch::kFloat32).device(vertices.device())),
             torch::empty({0, 8}, torch::TensorOptions().dtype(torch::kInt32).device(vertices.device())),
-            torch::empty({0}, torch::TensorOptions().dtype(torch::kInt64).device(vertices.device()))
+            return_unique_vert_ids ? std::make_optional(torch::empty({0}, torch::TensorOptions().dtype(torch::kInt64).device(vertices.device()))) : std::nullopt
+        );
+    }
+
+    auto vi = active_voxel_ids.div((ry - 1) * (rz - 1), "trunc");
+    auto rem = active_voxel_ids.remainder((ry - 1) * (rz - 1));
+    auto vj = rem.div(rz - 1, "trunc");
+    auto vk = rem.remainder(rz - 1);
+
+    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 active_voxels = torch::stack({v0, v1, v2, v3, v4, v5, v6, v7}, -1);
+
+    torch::Tensor unique_vert_ids, inverse_indices, counts;
+    std::tie(unique_vert_ids, inverse_indices, counts) = torch::_unique2(active_voxels.flatten(), true, true, false);
+
+    auto remapped_voxels = inverse_indices.view({-1, 8}).to(torch::kInt32);
+
+    auto u_i = unique_vert_ids.div(ry * rz, "trunc");
+    auto u_rem = unique_vert_ids.remainder(ry * rz);
+    auto u_j = u_rem.div(rz, "trunc");
+    auto u_k = u_rem.remainder(rz);
+
+    float spacing_x = (grid_max[0] - grid_min[0]) / (rx - 1);
+    float spacing_y = (grid_max[1] - grid_min[1]) / (ry - 1);
+    float spacing_z = (grid_max[2] - grid_min[2]) / (rz - 1);
+
+    auto x = grid_min[0] + u_i.to(torch::kFloat32) * spacing_x;
+    auto y = grid_min[1] + u_j.to(torch::kFloat32) * spacing_y;
+    auto z = grid_min[2] + u_k.to(torch::kFloat32) * spacing_z;
+
+    auto sparse_grid_vertices = torch::stack({x, y, z}, 1).contiguous();
+
+    return std::make_tuple(sparse_grid_vertices, remapped_voxels, return_unique_vert_ids ? std::make_optional(unique_vert_ids) : std::nullopt);
+}
+
+torch::Tensor get_active_voxel_ids_from_depth_py(
+    torch::Tensor depth_image,
+    torch::Tensor c2w_tensor,
+    torch::Tensor intrinsics_inv_tensor,
+    std::vector<float> grid_min,
+    std::vector<float> grid_max,
+    std::vector<int64_t> res,
+    bool activate_neighbor = false,
+    float trunc_margin = 0.0f
+) {
+    TORCH_CHECK(depth_image.is_cuda(), "depth_image must be a CUDA tensor");
+    TORCH_CHECK(depth_image.dim() == 2, "depth_image must be 2D (H, W)");
+    
+    int image_height = depth_image.size(0);
+    int image_width = depth_image.size(1);
+    int num_pixels = image_height * image_width;
+    
+    float4x4 c2w;
+    auto c2w_cpu = c2w_tensor.to(torch::kFloat32).cpu().contiguous();
+    auto c2w_a = c2w_cpu.accessor<float, 2>();
+    for(int r=0; r<4; ++r) for(int c=0; c<4; ++c) c2w.m[r][c] = c2w_a[r][c];
+    
+    float3x3 intrinsics_inv;
+    auto int_cpu = intrinsics_inv_tensor.to(torch::kFloat32).cpu().contiguous();
+    auto int_a = int_cpu.accessor<float, 2>();
+    for(int r=0; r<3; ++r) for(int c=0; c<3; ++c) intrinsics_inv.m[r][c] = int_a[r][c];
+    
+    float3 g_min; g_min.x = grid_min[0]; g_min.y = grid_min[1]; g_min.z = grid_min[2];
+    float3 g_max; g_max.x = grid_max[0]; g_max.y = grid_max[1]; g_max.z = grid_max[2];
+    int3 g_res; g_res.x = res[0]; g_res.y = res[1]; g_res.z = res[2];
+    
+    int max_pad_x = std::ceil(trunc_margin / ((grid_max[0] - grid_min[0]) / (res[0] - 1)));
+    int max_pad_y = std::ceil(trunc_margin / ((grid_max[1] - grid_min[1]) / (res[1] - 1)));
+    int max_pad_z = std::ceil(trunc_margin / ((grid_max[2] - grid_min[2]) / (res[2] - 1)));
+    int max_neighbors = activate_neighbor ? ((2 * max_pad_x + 1) * (2 * max_pad_y + 1) * (2 * max_pad_z + 1)) : 1;
+    
+    int64_t max_voxels = (int64_t)num_pixels * max_neighbors;
+    auto options = torch::TensorOptions().device(torch::kCUDA).dtype(torch::kInt64);
+    auto out_voxel_ids = torch::empty({max_voxels}, options);
+    
+    auto valid_counter_tensor = torch::zeros({1}, options);
+
+    grid::get_active_voxel_ids_from_depth(
+        num_pixels,
+        depth_image.data_ptr<float>(),
+        c2w,
+        intrinsics_inv,
+        image_width,
+        image_height,
+        g_min,
+        g_max,
+        g_res,
+        out_voxel_ids.data_ptr<int64_t>(),
+        reinterpret_cast<unsigned long long*>(valid_counter_tensor.data_ptr<int64_t>()),
+        activate_neighbor,
+        trunc_margin
+    );
+    
+    int64_t valid_count = valid_counter_tensor.item<int64_t>();
+    return out_voxel_ids.slice(0, 0, valid_count);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>> build_sparse_grid_from_active_voxels(
+    torch::Tensor active_voxel_ids,
+    std::vector<float> grid_min,
+    std::vector<float> grid_max,
+    std::vector<int64_t> res,
+    bool return_unique_vert_ids = true
+) {
+    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.");
+
+    int64_t rx = res[0];
+    int64_t ry = res[1];
+    int64_t rz = res[2];
+
+    if (active_voxel_ids.size(0) == 0) {
+        return std::make_tuple(
+            torch::empty({0, 3}, torch::TensorOptions().dtype(torch::kFloat32).device(active_voxel_ids.device())),
+            torch::empty({0, 8}, torch::TensorOptions().dtype(torch::kInt32).device(active_voxel_ids.device())),
+            return_unique_vert_ids ? std::make_optional(torch::empty({0}, torch::TensorOptions().dtype(torch::kInt64).device(active_voxel_ids.device()))) : std::nullopt
         );
     }
 
@@ -205,16 +336,20 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> create_voxel_grid_from_t
 
     auto sparse_grid_vertices = torch::stack({x, y, z}, 1).contiguous();
 
-    return std::make_tuple(sparse_grid_vertices, remapped_voxels, unique_vert_ids);
+    return std::make_tuple(sparse_grid_vertices, remapped_voxels, return_unique_vert_ids ? std::make_optional(unique_vert_ids) : std::nullopt);
 }
 
 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");
+          py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("device_str") = "cuda", py::arg("return_idx_grids") = true);
     m.def("create_voxel_grid_from_tmesh", &create_voxel_grid_from_tmesh, "Creates a sparse 3D voxel grid strictly around the surface.",
-          py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("tmesh"));
+          py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("tmesh"), py::arg("return_unique_vert_ids") = true);
     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"));
+    m.def("get_active_voxel_ids_from_depth", &get_active_voxel_ids_from_depth_py, "Extracts active voxel IDs from a single depth map.",
+          py::arg("depth_image"), py::arg("c2w_tensor"), py::arg("intrinsics_inv_tensor"), py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("activate_neighbor") = false, py::arg("trunc_margin") = 0.0f);
+    m.def("build_sparse_grid_from_active_voxels", &build_sparse_grid_from_active_voxels, "Builds sparse grid from unique voxel IDs.",
+          py::arg("active_voxel_ids"), py::arg("grid_min"), py::arg("grid_max"), py::arg("res"), py::arg("return_unique_vert_ids") = true);
 }

conquer3d-0.5.4/conquer3d/csrc/binds/ops/volint.cpp

@@ -0,0 +1,96 @@
+#include <torch/extension.h>
+#include <pybind11/pybind11.h>
+#include "../../ops/volint.h"
+#include "../../maths/maths.h"
+
+namespace py = pybind11;
+
+void single_view_volume_integral_bind(
+    torch::Tensor grid_vertices,
+    torch::Tensor sdf,
+    torch::Tensor weight,
+    std::optional<torch::Tensor> color_opt,
+    torch::Tensor depth_image,
+    std::optional<torch::Tensor> color_image_opt,
+    torch::Tensor extrinsics,
+    torch::Tensor intrinsics,
+    float trunc_margin,
+    int mode
+) {
+    TORCH_CHECK(grid_vertices.is_cuda() && grid_vertices.is_contiguous(), "grid_vertices must be CUDA and contiguous");
+    TORCH_CHECK(sdf.is_cuda() && sdf.is_contiguous(), "sdf must be CUDA and contiguous");
+    TORCH_CHECK(weight.is_cuda() && weight.is_contiguous(), "weight must be CUDA and contiguous");
+    TORCH_CHECK(depth_image.is_cuda() && depth_image.is_contiguous(), "depth_image must be CUDA and contiguous");
+
+    int num_vertices = grid_vertices.size(0);
+    int image_height = depth_image.size(0);
+    int image_width = depth_image.size(1);
+
+    const float3* grid_vertices_ptr = reinterpret_cast<const float3*>(grid_vertices.data_ptr<float>());
+    float* sdf_ptr = sdf.data_ptr<float>();
+    float* weight_ptr = weight.data_ptr<float>();
+    const float* depth_image_ptr = depth_image.data_ptr<float>();
+
+    float3* color_ptr = nullptr;
+    if (color_opt.has_value() && color_opt.value().defined() && color_opt.value().numel() > 0) {
+        auto color = color_opt.value();
+        TORCH_CHECK(color.is_cuda() && color.is_contiguous(), "color must be CUDA and contiguous");
+        color_ptr = reinterpret_cast<float3*>(color.data_ptr<float>());
+    }
+
+    const float3* color_image_ptr = nullptr;
+    if (color_image_opt.has_value() && color_image_opt.value().defined() && color_image_opt.value().numel() > 0) {
+        auto color_image = color_image_opt.value();
+        TORCH_CHECK(color_image.is_cuda() && color_image.is_contiguous(), "color_image must be CUDA and contiguous");
+        color_image_ptr = reinterpret_cast<const float3*>(color_image.data_ptr<float>());
+    }
+
+    torch::Tensor ex_cpu = extrinsics.cpu().contiguous();
+    torch::Tensor in_cpu = intrinsics.cpu().contiguous();
+    
+    const float* ex_data = ex_cpu.data_ptr<float>();
+    float4x4 ex = make_float4x4(
+        ex_data[0], ex_data[1], ex_data[2], ex_data[3],
+        ex_data[4], ex_data[5], ex_data[6], ex_data[7],
+        ex_data[8], ex_data[9], ex_data[10], ex_data[11],
+        ex_data[12], ex_data[13], ex_data[14], ex_data[15]
+    );
+
+    const float* in_data = in_cpu.data_ptr<float>();
+    float3x3 in = make_float3x3(
+        in_data[0], in_data[1], in_data[2],
+        in_data[3], in_data[4], in_data[5],
+        in_data[6], in_data[7], in_data[8]
+    );
+
+    single_view_volume_integral(
+        num_vertices,
+        grid_vertices_ptr,
+        sdf_ptr,
+        weight_ptr,
+        color_ptr,
+        depth_image_ptr,
+        color_image_ptr,
+        image_width,
+        image_height,
+        ex,
+        in,
+        trunc_margin,
+        mode
+    );
+}
+
+void bind_ops_volint(py::module_& m) {
+    m.def("single_view_volume_integral", &single_view_volume_integral_bind, 
+          py::arg("grid_vertices"),
+          py::arg("sdf"),
+          py::arg("weight"),
+          py::arg("color"),
+          py::arg("depth_image"),
+          py::arg("color_image"),
+          py::arg("extrinsics"),
+          py::arg("intrinsics"),
+          py::arg("trunc_margin"),
+          py::arg("mode") = 1,
+          "Single view TSDF volume integration");
+}

conquer3d-0.5.4/conquer3d/csrc/data_structure/grid.cu

@@ -0,0 +1,109 @@
+#include "grid.h"
+#include <stdio.h>
+#include <math.h>
+
+namespace grid {
+
+    __global__ void get_active_voxel_ids_from_depth_kernel(
+        const int num_pixels,
+        const float* depth_image,
+        const float4x4 c2w,
+        const float3x3 intrinsics_inv,
+        const int image_width,
+        const int image_height,
+        const float3 grid_min,
+        const float3 grid_max,
+        const int3 res,
+        int64_t* out_voxel_ids,
+        unsigned long long* valid_counter,
+        bool activate_neighbor,
+        float trunc_margin
+    ) {
+        int idx = blockIdx.x * blockDim.x + threadIdx.x;
+        if (idx >= num_pixels) return;
+
+        float d = depth_image[idx];
+        if (d <= 0.0f) return; // Ignore invalid depths
+
+        int ui = idx % image_width;
+        int vi = idx / image_width;
+
+        // 1. Unproject to camera coordinates using inverse intrinsics
+        float3 p_cam;
+        p_cam.x = d * (intrinsics_inv.m[0][0] * ui + intrinsics_inv.m[0][1] * vi + intrinsics_inv.m[0][2]);
+        p_cam.y = d * (intrinsics_inv.m[1][0] * ui + intrinsics_inv.m[1][1] * vi + intrinsics_inv.m[1][2]);
+        p_cam.z = d;
+
+        // 2. Transform to world coordinates using Camera-To-World matrix
+        float4 p_cam4 = make_float4(p_cam.x, p_cam.y, p_cam.z, 1.0f);
+        float4 p_world4 = c2w * p_cam4;
+        
+        // 3. Compute 3D grid cell indices (0-indexed)
+        float spacing_x = (grid_max.x - grid_min.x) / fmaxf(1.0f, (float)(res.x - 1));
+        float spacing_y = (grid_max.y - grid_min.y) / fmaxf(1.0f, (float)(res.y - 1));
+        float spacing_z = (grid_max.z - grid_min.z) / fmaxf(1.0f, (float)(res.z - 1));
+
+        int i_min = floorf((p_world4.x - grid_min.x) / spacing_x);
+        int i_max = i_min;
+        int j_min = floorf((p_world4.y - grid_min.y) / spacing_y);
+        int j_max = j_min;
+        int k_min = floorf((p_world4.z - grid_min.z) / spacing_z);
+        int k_max = k_min;
+
+        if (activate_neighbor) {
+            i_min = floorf((p_world4.x - trunc_margin - grid_min.x) / spacing_x);
+            i_max = floorf((p_world4.x + trunc_margin - grid_min.x) / spacing_x);
+            j_min = floorf((p_world4.y - trunc_margin - grid_min.y) / spacing_y);
+            j_max = floorf((p_world4.y + trunc_margin - grid_min.y) / spacing_y);
+            k_min = floorf((p_world4.z - trunc_margin - grid_min.z) / spacing_z);
+            k_max = floorf((p_world4.z + trunc_margin - grid_min.z) / spacing_z);
+        }
+
+        for (int i = i_min; i <= i_max; ++i) {
+            for (int j = j_min; j <= j_max; ++j) {
+                for (int k = k_min; k <= k_max; ++k) {
+                    if (i >= 0 && i < res.x - 1 && j >= 0 && j < res.y - 1 && k >= 0 && k < res.z - 1) {
+                        int64_t voxel_id = (int64_t)i * (res.y - 1) * (res.z - 1) + (int64_t)j * (res.z - 1) + (int64_t)k;
+                        unsigned long long write_idx = atomicAdd(valid_counter, 1ULL);
+                        out_voxel_ids[write_idx] = voxel_id;
+                    }
+                }
+            }
+        }
+    }
+
+    void get_active_voxel_ids_from_depth(
+        const int num_pixels,
+        const float* depth_image,
+        const float4x4 c2w,
+        const float3x3 intrinsics_inv,
+        const int image_width,
+        const int image_height,
+        const float3 grid_min,
+        const float3 grid_max,
+        const int3 res,
+        int64_t* out_voxel_ids,
+        unsigned long long* valid_counter,
+        bool activate_neighbor,
+        float trunc_margin
+    ) {
+        int block_size = 256;
+        int grid_size = (num_pixels + block_size - 1) / block_size;
+
+        get_active_voxel_ids_from_depth_kernel<<<grid_size, block_size>>>(
+            num_pixels,
+            depth_image,
+            c2w,
+            intrinsics_inv,
+            image_width,
+            image_height,
+            grid_min,
+            grid_max,
+            res,
+            out_voxel_ids,
+            valid_counter,
+            activate_neighbor,
+            trunc_margin
+        );
+    }
+}

conquer3d-0.5.4/conquer3d/csrc/data_structure/grid.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <vector>
+#include <torch/extension.h>
+#include "../maths/maths.h"
+
+namespace grid {
+    void get_active_voxel_ids_from_depth(
+        const int num_pixels,
+        const float* depth_image,
+        const float4x4 c2w,
+        const float3x3 intrinsics_inv,
+        const int image_width,
+        const int image_height,
+        const float3 grid_min,
+        const float3 grid_max,
+        const int3 res,
+        int64_t* out_voxel_ids,
+        unsigned long long* valid_counter,
+        bool activate_neighbor,
+        float trunc_margin
+    );
+}

{conquer3d-0.5.2 → conquer3d-0.5.4}/conquer3d/csrc/maths/f4x4.h

@@ -66,6 +66,7 @@ static __inline__ __host__ __device__ float4 operator* (const float4& a, const f
     );
 }
 
+// [4, 4] x [4, 1] = [4, 1]
 static __inline__ __host__ __device__ float4 operator* (const float4x4& m, const float4& a)
 {
     return make_float4(