conquer3d 0.3.0__tar.gz → 0.3.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {conquer3d-0.3.0/conquer3d.egg-info → conquer3d-0.3.2}/PKG-INFO +5 -2
- {conquer3d-0.3.0 → conquer3d-0.3.2}/README.md +3 -1
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/_C.so +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/__init__.py +2 -1
- conquer3d-0.3.2/conquer3d/conversion/__init__.py +13 -0
- conquer3d-0.3.2/conquer3d/conversion/grid.py +156 -0
- conquer3d-0.3.2/conquer3d/conversion/mesh.py +120 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/grid.cpp +1 -2
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/aabb.h +15 -0
- conquer3d-0.3.2/conquer3d/data/collate/__init__.py +0 -0
- conquer3d-0.3.2/conquer3d/data/collate/mesh.py +40 -0
- conquer3d-0.3.2/conquer3d/data/collate/sparse_tensor.py +43 -0
- conquer3d-0.3.2/conquer3d/data/dataset/__init__.py +0 -0
- conquer3d-0.3.2/conquer3d/data/dataset/base_mesh.py +17 -0
- conquer3d-0.3.2/conquer3d/data/dataset/digit3d.py +132 -0
- conquer3d-0.3.2/conquer3d/data/transform/__init__.py +0 -0
- conquer3d-0.3.2/conquer3d/data/transform/base.py +39 -0
- conquer3d-0.3.2/conquer3d/data/transform/ops.py +62 -0
- conquer3d-0.3.2/conquer3d/data/transform/vertex.py +55 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/data_structure/__init__.py +10 -2
- conquer3d-0.3.2/conquer3d/data_structure/bmesh.py +23 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/data_structure/grid.py +1 -50
- conquer3d-0.3.2/conquer3d/primitive/__init__.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2/conquer3d.egg-info}/PKG-INFO +5 -2
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d.egg-info/SOURCES.txt +15 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d.egg-info/requires.txt +1 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/pyproject.toml +2 -1
- {conquer3d-0.3.0 → conquer3d-0.3.2}/LICENSE +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/MANIFEST.in +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/_C.pyi +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/creation/__init__.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/creation/triangle_creation.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/creation/triangle_creation.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/bvh.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/gs_bvh.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/kdtree.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/mesh_bvh.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/pgs_bvh.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/data_structure/triangle_mesh.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/ops/chamfer.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/ops/mc.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/primitive/gs.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/binds/primitive/pgs.cpp +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/check.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/constants.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/creation/triangle_creation.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/bvh.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/bvh.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/gs_bvh.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/kdtree.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/kdtree.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/mesh_bvh.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/mesh_bvh.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/pgs_bvh.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/triangle_mesh.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/triangle_mesh.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/data_structure/zcurve.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f2x2.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f3x1.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f3x3.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f3x4.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f4x1.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/f4x4.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/maths.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/maths/ops.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/ops/chamfer.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/ops/chamfer.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/ops/mc.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/ops/mc.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/ops/mc_data.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/edge.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/gs.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/gs.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/gs_aabb.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/gs_math.cuh +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/pgs.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/pgs.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/pgs_aabb.cu +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/pgs_math.cuh +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/ray.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/primitive/triangles.h +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/csrc/pybind.cpp +0 -0
- {conquer3d-0.3.0/conquer3d/primitive → conquer3d-0.3.2/conquer3d/data}/__init__.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/ops/__init__.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/ops/diff_marching_cubes.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/ops/distance.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/ops/marching_cubes.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/primitive/gs.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d/primitive/pgs.py +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d.egg-info/dependency_links.txt +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d.egg-info/not-zip-safe +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/conquer3d.egg-info/top_level.txt +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/setup.cfg +0 -0
- {conquer3d-0.3.0 → conquer3d-0.3.2}/setup.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: conquer3d
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: Geometric Cuda Tool Box
|
|
5
5
|
Author-email: Do Hoang Khoi <khoido8899@gmail.com>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -13,6 +13,7 @@ Requires-Python: >=3.9
|
|
|
13
13
|
Description-Content-Type: text/markdown
|
|
14
14
|
License-File: LICENSE
|
|
15
15
|
Requires-Dist: torch>=2.0.0
|
|
16
|
+
Requires-Dist: gdown
|
|
16
17
|
Provides-Extra: cuda
|
|
17
18
|
Requires-Dist: cuda-toolkit==12.8.2; extra == "cuda"
|
|
18
19
|
Provides-Extra: dev
|
|
@@ -61,7 +62,9 @@ pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/t
|
|
|
61
62
|
pip install pybind11-stubgen
|
|
62
63
|
pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
|
|
63
64
|
|
|
64
|
-
pip install plotly open3d jupyter trimesh point-cloud-utils
|
|
65
|
+
pip install plotly open3d jupyter trimesh point-cloud-utils meshlib
|
|
66
|
+
|
|
67
|
+
pip install rectified-flow-pytorch
|
|
65
68
|
```
|
|
66
69
|
|
|
67
70
|
# Development
|
|
@@ -33,7 +33,9 @@ pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/t
|
|
|
33
33
|
pip install pybind11-stubgen
|
|
34
34
|
pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
|
|
35
35
|
|
|
36
|
-
pip install plotly open3d jupyter trimesh point-cloud-utils
|
|
36
|
+
pip install plotly open3d jupyter trimesh point-cloud-utils meshlib
|
|
37
|
+
|
|
38
|
+
pip install rectified-flow-pytorch
|
|
37
39
|
```
|
|
38
40
|
|
|
39
41
|
# Development
|
|
Binary file
|
|
@@ -5,5 +5,6 @@ from . import creation
|
|
|
5
5
|
from . import data_structure
|
|
6
6
|
from . import primitive
|
|
7
7
|
from . import ops
|
|
8
|
+
from . import conversion
|
|
8
9
|
|
|
9
|
-
__all__ = ['_C', 'creation', 'data_structure', 'primitive', 'ops']
|
|
10
|
+
__all__ = ['_C', 'creation', 'data_structure', 'primitive', 'ops', 'conversion']
|
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from typing import Tuple, List, Union
|
|
3
|
+
|
|
4
|
+
def voxel2sparse(
|
|
5
|
+
active_indices: torch.Tensor,
|
|
6
|
+
voxels: torch.Tensor,
|
|
7
|
+
idx_grids: torch.Tensor,
|
|
8
|
+
sdf: torch.Tensor = None,
|
|
9
|
+
batch_idx: int = 0
|
|
10
|
+
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
|
|
11
|
+
"""
|
|
12
|
+
Converts 1D active voxel indices into the [N, 4] coordinates tensor
|
|
13
|
+
expected by sparse neural networks (like spconv/torchsparse).
|
|
14
|
+
|
|
15
|
+
Args:
|
|
16
|
+
active_indices (torch.Tensor): 1D tensor of active voxel indices.
|
|
17
|
+
voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices for each voxel.
|
|
18
|
+
idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) 3D coordinate indices for each vertex.
|
|
19
|
+
sdf (torch.Tensor, optional): A float32 tensor of shape (N,) containing the SDF values. If provided, the function will also extract the 8 corner SDFs for each active voxel. Defaults to None.
|
|
20
|
+
batch_idx (int, optional): The batch index to prepend. Defaults to 0.
|
|
21
|
+
|
|
22
|
+
Returns:
|
|
23
|
+
If sdf is None:
|
|
24
|
+
torch.Tensor: An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
|
|
25
|
+
If sdf is provided:
|
|
26
|
+
Tuple[torch.Tensor, torch.Tensor]:
|
|
27
|
+
- sparse_coords (torch.Tensor): [N, 4] tensor
|
|
28
|
+
- sparse_sdfs (torch.Tensor): [N, 8] tensor containing the 8 corner SDF values for each voxel.
|
|
29
|
+
"""
|
|
30
|
+
active_corners = voxels[active_indices]
|
|
31
|
+
base_vertex_indices = active_corners[:, 0].to(torch.int64)
|
|
32
|
+
voxel_coords = idx_grids[base_vertex_indices]
|
|
33
|
+
|
|
34
|
+
batch_col = torch.full((voxel_coords.shape[0], 1), batch_idx, device=voxel_coords.device, dtype=torch.int64)
|
|
35
|
+
sparse_coords = torch.cat([batch_col, voxel_coords], dim=1).to(torch.int32)
|
|
36
|
+
|
|
37
|
+
if sdf is not None:
|
|
38
|
+
sparse_sdfs = sdf[active_corners.to(torch.int64)]
|
|
39
|
+
return sparse_coords, sparse_sdfs
|
|
40
|
+
|
|
41
|
+
return sparse_coords
|
|
42
|
+
|
|
43
|
+
def sparse2voxel(
|
|
44
|
+
sparse_coords: torch.Tensor,
|
|
45
|
+
sparse_sdfs: torch.Tensor,
|
|
46
|
+
grid_min: Union[List[float], Tuple[float, float, float]],
|
|
47
|
+
grid_max: Union[List[float], Tuple[float, float, float]],
|
|
48
|
+
res: Union[List[int], Tuple[int, int, int]]
|
|
49
|
+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
50
|
+
"""
|
|
51
|
+
Converts pure voxel-centric sparse representations back into a localized,
|
|
52
|
+
dense-free mesh topology ready for marching cubes.
|
|
53
|
+
|
|
54
|
+
Args:
|
|
55
|
+
sparse_coords (torch.Tensor): [N, 4] tensor of voxel coords [batch_idx, x, y, z].
|
|
56
|
+
sparse_sdfs (torch.Tensor): [N, 8] tensor of corner SDFs.
|
|
57
|
+
grid_min (List[float]): Global minimum (x, y, z) bounds.
|
|
58
|
+
grid_max (List[float]): Global maximum (x, y, z) bounds.
|
|
59
|
+
res (List[int]): Global resolution (rx, ry, rz).
|
|
60
|
+
|
|
61
|
+
Returns:
|
|
62
|
+
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
63
|
+
- unique_vertices (torch.Tensor): [M, 3] float coordinates of all unique mesh vertices.
|
|
64
|
+
- local_voxels (torch.Tensor): [N, 8] int tensor of corner indices mapping to unique_vertices.
|
|
65
|
+
- merged_sdfs (torch.Tensor): [M,] float tensor of averaged SDF values at unique_vertices.
|
|
66
|
+
"""
|
|
67
|
+
device = sparse_coords.device
|
|
68
|
+
|
|
69
|
+
# 1. Base integer coordinates of the voxels
|
|
70
|
+
base_coords = sparse_coords[:, 1:] # [N, 3]
|
|
71
|
+
|
|
72
|
+
# 2. Local offsets for the 8 corners of a cube
|
|
73
|
+
cube_corners = torch.tensor([
|
|
74
|
+
[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0],
|
|
75
|
+
[0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]
|
|
76
|
+
], dtype=torch.int32, device=device)
|
|
77
|
+
|
|
78
|
+
# 3. Compute all corner integer coordinates [N, 8, 3]
|
|
79
|
+
all_corners = base_coords.unsqueeze(1) + cube_corners.unsqueeze(0)
|
|
80
|
+
all_corners_flat = all_corners.reshape(-1, 3) # [N*8, 3]
|
|
81
|
+
|
|
82
|
+
# 4. Find the unique vertices to form our local mesh graph
|
|
83
|
+
unique_corners_int, inverse_indices = torch.unique(all_corners_flat, dim=0, return_inverse=True)
|
|
84
|
+
local_voxels = inverse_indices.reshape(-1, 8).to(torch.int32) # [N, 8]
|
|
85
|
+
|
|
86
|
+
# 5. Average the SDFs at the overlapping corners
|
|
87
|
+
num_unique = unique_corners_int.shape[0]
|
|
88
|
+
merged_sdfs = torch.zeros(num_unique, dtype=sparse_sdfs.dtype, device=device)
|
|
89
|
+
counts = torch.zeros(num_unique, dtype=sparse_sdfs.dtype, device=device)
|
|
90
|
+
|
|
91
|
+
flat_inverse = inverse_indices # [N*8]
|
|
92
|
+
flat_sdfs = sparse_sdfs.flatten() # [N*8]
|
|
93
|
+
|
|
94
|
+
merged_sdfs.scatter_add_(0, flat_inverse, flat_sdfs)
|
|
95
|
+
counts.scatter_add_(0, flat_inverse, torch.ones_like(flat_sdfs))
|
|
96
|
+
merged_sdfs = merged_sdfs / counts
|
|
97
|
+
|
|
98
|
+
# 6. Convert integer coordinates back to global float positions
|
|
99
|
+
rx, ry, rz = res
|
|
100
|
+
x_coords = torch.linspace(grid_min[0], grid_max[0], rx + 1, device=device, dtype=torch.float32)
|
|
101
|
+
y_coords = torch.linspace(grid_min[1], grid_max[1], ry + 1, device=device, dtype=torch.float32)
|
|
102
|
+
z_coords = torch.linspace(grid_min[2], grid_max[2], rz + 1, device=device, dtype=torch.float32)
|
|
103
|
+
|
|
104
|
+
ux = x_coords[unique_corners_int[:, 0]]
|
|
105
|
+
uy = y_coords[unique_corners_int[:, 1]]
|
|
106
|
+
uz = z_coords[unique_corners_int[:, 2]]
|
|
107
|
+
|
|
108
|
+
unique_vertices = torch.stack([ux, uy, uz], dim=1)
|
|
109
|
+
|
|
110
|
+
return unique_vertices, local_voxels, merged_sdfs
|
|
111
|
+
|
|
112
|
+
def sparse_coo2dense_occ(
|
|
113
|
+
sparse_coords: torch.Tensor,
|
|
114
|
+
batch_size: int,
|
|
115
|
+
res: Union[List[int], Tuple[int, int, int]]
|
|
116
|
+
) -> torch.Tensor:
|
|
117
|
+
"""
|
|
118
|
+
Converts a [N, 4] sparse coordinates tensor into a [B, 1, rx, ry, rz] dense occupancy grid.
|
|
119
|
+
|
|
120
|
+
Args:
|
|
121
|
+
sparse_coords (torch.Tensor): An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
|
|
122
|
+
batch_size (int): The batch size B.
|
|
123
|
+
res (List[int] | Tuple[int, int, int]): The grid resolution (rx, ry, rz).
|
|
124
|
+
|
|
125
|
+
Returns:
|
|
126
|
+
torch.Tensor: A float32 tensor of shape [B, 1, rx, ry, rz] containing 1.0 at active voxels and 0.0 otherwise.
|
|
127
|
+
"""
|
|
128
|
+
rx, ry, rz = res
|
|
129
|
+
dense_occ = torch.zeros((batch_size, 1, rx, ry, rz), dtype=torch.float32, device=sparse_coords.device)
|
|
130
|
+
|
|
131
|
+
b = sparse_coords[:, 0].to(torch.int64)
|
|
132
|
+
x = sparse_coords[:, 1].to(torch.int64)
|
|
133
|
+
y = sparse_coords[:, 2].to(torch.int64)
|
|
134
|
+
z = sparse_coords[:, 3].to(torch.int64)
|
|
135
|
+
|
|
136
|
+
dense_occ[b, 0, x, y, z] = 1.0
|
|
137
|
+
return dense_occ
|
|
138
|
+
|
|
139
|
+
def dense_occ2sparse_coo(
|
|
140
|
+
dense_occ: torch.Tensor,
|
|
141
|
+
threshold: float = 0.0
|
|
142
|
+
) -> torch.Tensor:
|
|
143
|
+
"""
|
|
144
|
+
Converts a [B, C, rx, ry, rz] dense occupancy grid back to a [N, 4] sparse coordinates tensor.
|
|
145
|
+
|
|
146
|
+
Args:
|
|
147
|
+
dense_occ (torch.Tensor): A tensor of shape [B, C, rx, ry, rz].
|
|
148
|
+
|
|
149
|
+
Returns:
|
|
150
|
+
torch.Tensor: An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
|
|
151
|
+
"""
|
|
152
|
+
active_mask = dense_occ > threshold
|
|
153
|
+
indices = torch.argwhere(active_mask)
|
|
154
|
+
# indices has shape [N, 5] for [b, c, x, y, z]. We want [batch_idx, x, y, z]
|
|
155
|
+
sparse_coords = indices[:, [0, 2, 3, 4]].to(torch.int32)
|
|
156
|
+
return sparse_coords
|
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from conquer3d.data_structure.grid import create_voxel_grid, compute_active_voxels
|
|
3
|
+
from conquer3d.conversion.grid import voxel2sparse
|
|
4
|
+
from conquer3d.data_structure import TriangleMesh
|
|
5
|
+
from conquer3d.data_structure.bmesh import BTriangleMesh
|
|
6
|
+
from typing import Union, List, Tuple
|
|
7
|
+
|
|
8
|
+
def mesh2voxel(vertices: torch.Tensor, res: Union[List[int], Tuple[int, int, int]], grid_bound: float = None):
|
|
9
|
+
"""
|
|
10
|
+
Computes bounding box of mesh and returns a structured voxel grid.
|
|
11
|
+
"""
|
|
12
|
+
if grid_bound is None:
|
|
13
|
+
grid_min = vertices.min(dim=0)[0].tolist()
|
|
14
|
+
grid_max = vertices.max(dim=0)[0].tolist()
|
|
15
|
+
else:
|
|
16
|
+
grid_min = [-grid_bound, -grid_bound, -grid_bound]
|
|
17
|
+
grid_max = [grid_bound, grid_bound, grid_bound]
|
|
18
|
+
|
|
19
|
+
device_str = str(vertices.device)
|
|
20
|
+
if device_str.startswith('cuda'):
|
|
21
|
+
# Fix format if there's an index like cuda:0
|
|
22
|
+
device_str = "cuda"
|
|
23
|
+
|
|
24
|
+
grid_vertices, voxels, idx_grids = create_voxel_grid(grid_min, grid_max, res, device=device_str)
|
|
25
|
+
|
|
26
|
+
# Ensure they are on the exact same device as vertices (especially if multiple GPUs)
|
|
27
|
+
grid_vertices = grid_vertices.to(vertices.device)
|
|
28
|
+
voxels = voxels.to(vertices.device)
|
|
29
|
+
idx_grids = idx_grids.to(vertices.device)
|
|
30
|
+
|
|
31
|
+
return grid_vertices, voxels, idx_grids, grid_min, grid_max
|
|
32
|
+
|
|
33
|
+
def mesh2voxel_sdf(vertices: torch.Tensor, faces: torch.Tensor, grid_vertices: torch.Tensor):
|
|
34
|
+
"""
|
|
35
|
+
Computes SDF for grid vertices against the mesh using C++ TriangleMesh fast winding numbers.
|
|
36
|
+
"""
|
|
37
|
+
mesh = TriangleMesh(vertices, faces.to(torch.int32))
|
|
38
|
+
|
|
39
|
+
# query_points returns: q_ids, tri_ids, prj_pts, sdf
|
|
40
|
+
_, _, _, sdf = mesh.query_points(
|
|
41
|
+
grid_vertices,
|
|
42
|
+
return_sdf=True,
|
|
43
|
+
return_prj_pts=False,
|
|
44
|
+
sign_mode=0,
|
|
45
|
+
distance_mode=0
|
|
46
|
+
)
|
|
47
|
+
return sdf
|
|
48
|
+
|
|
49
|
+
def mesh2sparse(bmesh: BTriangleMesh, res: Union[List[int], Tuple[int, int, int]], grid_bound: float = 1.2, iso: float = 0.0):
|
|
50
|
+
"""
|
|
51
|
+
Batched processing of BTriangleMesh into sparse tensors.
|
|
52
|
+
"""
|
|
53
|
+
all_coords = []
|
|
54
|
+
all_sdfs = []
|
|
55
|
+
|
|
56
|
+
for b in range(bmesh.batch_size):
|
|
57
|
+
# Extract individual mesh
|
|
58
|
+
v_mask = bmesh.vertbids == b
|
|
59
|
+
f_mask = bmesh.facebids == b
|
|
60
|
+
|
|
61
|
+
v = bmesh.vertices[v_mask]
|
|
62
|
+
f = bmesh.faces[f_mask]
|
|
63
|
+
|
|
64
|
+
if len(v) == 0:
|
|
65
|
+
continue
|
|
66
|
+
|
|
67
|
+
# Voxelize and Compute SDF
|
|
68
|
+
grid_vertices, voxels, idx_grids, _, _ = mesh2voxel(v, res, grid_bound)
|
|
69
|
+
sdf = mesh2voxel_sdf(v, f, grid_vertices)
|
|
70
|
+
|
|
71
|
+
# Sparse conversion
|
|
72
|
+
active_voxels = compute_active_voxels(voxels, sdf, iso)
|
|
73
|
+
sparse_coords, sparse_sdfs = voxel2sparse(active_voxels, voxels, idx_grids, sdf=sdf, batch_idx=b)
|
|
74
|
+
|
|
75
|
+
all_coords.append(sparse_coords)
|
|
76
|
+
all_sdfs.append(sparse_sdfs)
|
|
77
|
+
|
|
78
|
+
if len(all_coords) > 0:
|
|
79
|
+
return torch.cat(all_coords, dim=0), torch.cat(all_sdfs, dim=0)
|
|
80
|
+
else:
|
|
81
|
+
# Fallback for empty batch
|
|
82
|
+
device = bmesh.vertices.device
|
|
83
|
+
return torch.empty((0, 4), dtype=torch.int32, device=device), torch.empty((0, 8), dtype=torch.float32, device=device)
|
|
84
|
+
|
|
85
|
+
def mesh2sparse_with_dense(bmesh: BTriangleMesh, res: Union[List[int], Tuple[int, int, int]], grid_bound: float = 1.2, iso: float = 0.0):
|
|
86
|
+
"""
|
|
87
|
+
Batched processing of BTriangleMesh into sparse tensors, also returning full dense coords.
|
|
88
|
+
"""
|
|
89
|
+
all_coords = []
|
|
90
|
+
all_sdfs = []
|
|
91
|
+
all_dense_coords = []
|
|
92
|
+
all_dense_sdfs = []
|
|
93
|
+
|
|
94
|
+
for b in range(bmesh.batch_size):
|
|
95
|
+
v_mask = bmesh.vertbids == b
|
|
96
|
+
f_mask = bmesh.facebids == b
|
|
97
|
+
|
|
98
|
+
v = bmesh.vertices[v_mask]
|
|
99
|
+
f = bmesh.faces[f_mask]
|
|
100
|
+
|
|
101
|
+
if len(v) == 0:
|
|
102
|
+
continue
|
|
103
|
+
|
|
104
|
+
grid_vertices, voxels, idx_grids, _, _ = mesh2voxel(v, res, grid_bound)
|
|
105
|
+
sdf = mesh2voxel_sdf(v, f, grid_vertices)
|
|
106
|
+
|
|
107
|
+
active_voxels = compute_active_voxels(voxels, sdf, iso)
|
|
108
|
+
sparse_coords, sparse_sdfs = voxel2sparse(active_voxels, voxels, idx_grids, sdf=sdf, batch_idx=b)
|
|
109
|
+
dense_coords, dense_sdfs = voxel2sparse(torch.ones_like(voxels[:, 0], dtype=torch.bool), voxels, idx_grids, sdf=sdf, batch_idx=b)
|
|
110
|
+
|
|
111
|
+
all_coords.append(sparse_coords)
|
|
112
|
+
all_sdfs.append(sparse_sdfs)
|
|
113
|
+
all_dense_coords.append(dense_coords)
|
|
114
|
+
all_dense_sdfs.append(dense_sdfs)
|
|
115
|
+
|
|
116
|
+
if len(all_coords) > 0:
|
|
117
|
+
return torch.cat(all_coords, dim=0), torch.cat(all_sdfs, dim=0), torch.cat(all_dense_coords, dim=0), torch.cat(all_dense_sdfs, dim=0)
|
|
118
|
+
else:
|
|
119
|
+
device = bmesh.vertices.device
|
|
120
|
+
return torch.empty((0, 4), dtype=torch.int32, device=device), torch.empty((0, 8), dtype=torch.float32, device=device), torch.empty((0, 4), dtype=torch.int32, device=device), torch.empty((0, 8), dtype=torch.float32, device=device)
|
|
@@ -122,8 +122,7 @@ torch::Tensor compute_grid_normal(torch::Tensor sdf, torch::Tensor grid_vertices
|
|
|
122
122
|
}
|
|
123
123
|
|
|
124
124
|
torch::Tensor compute_active_voxels(torch::Tensor voxels, torch::Tensor sdf, float iso) {
|
|
125
|
-
TORCH_CHECK(voxels.
|
|
126
|
-
TORCH_CHECK(sdf.is_cuda(), "sdf must be a CUDA tensor.");
|
|
125
|
+
TORCH_CHECK(voxels.device() == sdf.device(), "voxels and sdf must be on the same device.");
|
|
127
126
|
TORCH_CHECK(voxels.dtype() == torch::kInt32, "voxels must be int32.");
|
|
128
127
|
TORCH_CHECK(sdf.dtype() == torch::kFloat32, "sdf must be float32.");
|
|
129
128
|
|
|
@@ -167,6 +167,21 @@ namespace aabb
|
|
|
167
167
|
(target_aabb_max.y) >= query_aabb_max.y &&
|
|
168
168
|
(target_aabb_max.z) >= query_aabb_max.z);
|
|
169
169
|
}
|
|
170
|
+
|
|
171
|
+
__device__ __forceinline__ bool test_point_inside_aabb(
|
|
172
|
+
const float3 &point,
|
|
173
|
+
const float3 &aabb_min,
|
|
174
|
+
const float3 &aabb_max
|
|
175
|
+
)
|
|
176
|
+
{
|
|
177
|
+
return (
|
|
178
|
+
aabb_min.x <= point.x &&
|
|
179
|
+
aabb_min.y <= point.y &&
|
|
180
|
+
aabb_min.z <= point.z &&
|
|
181
|
+
(aabb_max.x) >= point.x &&
|
|
182
|
+
(aabb_max.y) >= point.y &&
|
|
183
|
+
(aabb_max.z) >= point.z);
|
|
184
|
+
}
|
|
170
185
|
}
|
|
171
186
|
|
|
172
187
|
#endif // AABB_H
|
|
File without changes
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from conquer3d.data_structure.bmesh import BTriangleMesh
|
|
3
|
+
|
|
4
|
+
def bmesh_collate_fn(batch):
|
|
5
|
+
"""
|
|
6
|
+
Collate function for a dataset returning (vertices, faces, label).
|
|
7
|
+
Returns a BTriangleMesh and a batched label tensor.
|
|
8
|
+
"""
|
|
9
|
+
all_vertices = []
|
|
10
|
+
all_faces = []
|
|
11
|
+
all_vertbids = []
|
|
12
|
+
all_facebids = []
|
|
13
|
+
all_labels = []
|
|
14
|
+
|
|
15
|
+
for b, (v, f, l) in enumerate(batch):
|
|
16
|
+
all_vertices.append(v)
|
|
17
|
+
|
|
18
|
+
# We do not offset the face indices because we will process them per-mesh locally
|
|
19
|
+
all_faces.append(f)
|
|
20
|
+
|
|
21
|
+
all_vertbids.append(torch.full((v.shape[0],), b, dtype=torch.int32))
|
|
22
|
+
all_facebids.append(torch.full((f.shape[0],), b, dtype=torch.int32))
|
|
23
|
+
all_labels.append(l)
|
|
24
|
+
|
|
25
|
+
batched_vertices = torch.cat(all_vertices, dim=0)
|
|
26
|
+
batched_faces = torch.cat(all_faces, dim=0)
|
|
27
|
+
batched_vertbids = torch.cat(all_vertbids, dim=0)
|
|
28
|
+
batched_facebids = torch.cat(all_facebids, dim=0)
|
|
29
|
+
|
|
30
|
+
batched_labels = torch.tensor(all_labels, dtype=torch.long)
|
|
31
|
+
|
|
32
|
+
bmesh = BTriangleMesh(
|
|
33
|
+
vertices=batched_vertices,
|
|
34
|
+
faces=batched_faces,
|
|
35
|
+
vertbids=batched_vertbids,
|
|
36
|
+
facebids=batched_facebids,
|
|
37
|
+
batch_size=len(batch)
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
return bmesh, batched_labels
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
|
|
3
|
+
def sparse_collate_fn(batch):
|
|
4
|
+
"""
|
|
5
|
+
Collate function for sparse tensor dataloaders (e.g., MinkowskiEngine, spconv, torchsparse).
|
|
6
|
+
Takes a batch of samples where each sample is (idx_grids, features, ...).
|
|
7
|
+
Appends the batch index to the idx_grids to create (batch_idx, x, y, z).
|
|
8
|
+
|
|
9
|
+
Args:
|
|
10
|
+
batch (list): A list of tuples (idx_grids, features, label)
|
|
11
|
+
|
|
12
|
+
Returns:
|
|
13
|
+
tuple: (batched_coords, batched_features, batched_labels)
|
|
14
|
+
batched_coords (torch.Tensor): Shape (Total_N, 4)
|
|
15
|
+
batched_features (torch.Tensor): Shape (Total_N, ...)
|
|
16
|
+
batched_labels (torch.Tensor): Shape (Batch_Size,)
|
|
17
|
+
"""
|
|
18
|
+
batched_coords = []
|
|
19
|
+
batched_features = []
|
|
20
|
+
batched_labels = []
|
|
21
|
+
|
|
22
|
+
for batch_idx, item in enumerate(batch):
|
|
23
|
+
idx_grids, features, label = item
|
|
24
|
+
|
|
25
|
+
# Create a batch index column of shape [N, 1]
|
|
26
|
+
b_col = torch.full((idx_grids.shape[0], 1), batch_idx, dtype=idx_grids.dtype, device=idx_grids.device)
|
|
27
|
+
|
|
28
|
+
# Concatenate to get [N, 4] formatted as (batch_idx, x, y, z)
|
|
29
|
+
coords_with_batch = torch.cat([b_col, idx_grids], dim=1)
|
|
30
|
+
|
|
31
|
+
batched_coords.append(coords_with_batch)
|
|
32
|
+
batched_features.append(features)
|
|
33
|
+
batched_labels.append(label)
|
|
34
|
+
|
|
35
|
+
batched_coords = torch.cat(batched_coords, dim=0)
|
|
36
|
+
batched_features = torch.cat(batched_features, dim=0)
|
|
37
|
+
|
|
38
|
+
if isinstance(batched_labels[0], torch.Tensor):
|
|
39
|
+
batched_labels = torch.stack(batched_labels, dim=0)
|
|
40
|
+
else:
|
|
41
|
+
batched_labels = torch.tensor(batched_labels)
|
|
42
|
+
|
|
43
|
+
return batched_coords, batched_features, batched_labels
|
|
File without changes
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from torch.utils.data import Dataset
|
|
2
|
+
from typing import Callable, Optional
|
|
3
|
+
|
|
4
|
+
class BaseMeshDataset(Dataset):
|
|
5
|
+
"""
|
|
6
|
+
Base PyTorch Dataset class for Mesh datasets.
|
|
7
|
+
"""
|
|
8
|
+
def __init__(self, root: str, transform: Optional[Callable] = None):
|
|
9
|
+
super().__init__()
|
|
10
|
+
self.root = root
|
|
11
|
+
self.transform = transform
|
|
12
|
+
|
|
13
|
+
def __len__(self) -> int:
|
|
14
|
+
raise NotImplementedError("Subclasses must implement __len__")
|
|
15
|
+
|
|
16
|
+
def __getitem__(self, idx: int):
|
|
17
|
+
raise NotImplementedError("Subclasses must implement __getitem__")
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import zipfile
|
|
3
|
+
import torch
|
|
4
|
+
import conquer3d as c3d
|
|
5
|
+
import meshlib.mrmeshpy as mr
|
|
6
|
+
import meshlib.mrmeshnumpy as mrnp
|
|
7
|
+
|
|
8
|
+
from .base_mesh import BaseMeshDataset
|
|
9
|
+
|
|
10
|
+
class Digit3D(BaseMeshDataset):
|
|
11
|
+
"""
|
|
12
|
+
Digit3D Mesh Dataset containing 3D MNIST digits.
|
|
13
|
+
"""
|
|
14
|
+
def __init__(self, root: str = "~/.conquer3d/", train: bool = True, transform=None, download: bool = False, cached: bool = False):
|
|
15
|
+
root = os.path.expanduser(root)
|
|
16
|
+
super().__init__(root, transform)
|
|
17
|
+
self.train = train
|
|
18
|
+
self.zip_path = os.path.join(root, "digit3d.zip")
|
|
19
|
+
self.split_dir = "src/train" if train else "src/test"
|
|
20
|
+
self.cached = cached
|
|
21
|
+
self._cache = {}
|
|
22
|
+
|
|
23
|
+
if download:
|
|
24
|
+
self.download()
|
|
25
|
+
|
|
26
|
+
if not os.path.exists(self.zip_path):
|
|
27
|
+
raise RuntimeError(f"Dataset not found at {self.zip_path}. You can use download=True to download it.")
|
|
28
|
+
|
|
29
|
+
with zipfile.ZipFile(self.zip_path, 'r') as z:
|
|
30
|
+
self.all_files = [f for f in z.namelist() if f.startswith(self.split_dir) and f.endswith(".obj")]
|
|
31
|
+
|
|
32
|
+
def download(self):
|
|
33
|
+
if os.path.exists(self.zip_path):
|
|
34
|
+
return
|
|
35
|
+
os.makedirs(self.root, exist_ok=True)
|
|
36
|
+
url = "https://drive.google.com/uc?id=1Vry0-sflcSmpwZnjn8yBbF2vBfuW1T_W"
|
|
37
|
+
try:
|
|
38
|
+
import gdown
|
|
39
|
+
except ImportError:
|
|
40
|
+
raise ImportError("gdown is required to download the dataset. Please install it using 'pip install gdown'.")
|
|
41
|
+
print(f"Downloading Digit3D dataset to {self.zip_path}...")
|
|
42
|
+
gdown.download(url, self.zip_path, quiet=False)
|
|
43
|
+
|
|
44
|
+
def __len__(self) -> int:
|
|
45
|
+
return len(self.all_files)
|
|
46
|
+
|
|
47
|
+
def __getitem__(self, idx: int):
|
|
48
|
+
if self.cached and idx in self._cache:
|
|
49
|
+
vertices_t, faces_t, label = self._cache[idx]
|
|
50
|
+
if self.transform:
|
|
51
|
+
vertices_t, faces_t = self.transform(vertices_t.clone(), faces_t.clone())
|
|
52
|
+
return vertices_t, faces_t, label
|
|
53
|
+
|
|
54
|
+
f_path = self.all_files[idx]
|
|
55
|
+
basename = os.path.basename(f_path)
|
|
56
|
+
label = int(basename.split("_")[0])
|
|
57
|
+
|
|
58
|
+
# Read directly from zip stream to avoid file descriptor and extraction I/O overhead
|
|
59
|
+
if getattr(self, '_zip', None) is None:
|
|
60
|
+
self._zip = zipfile.ZipFile(self.zip_path, 'r')
|
|
61
|
+
|
|
62
|
+
with self._zip.open(f_path, 'r') as f:
|
|
63
|
+
content = f.read().decode('utf-8')
|
|
64
|
+
|
|
65
|
+
vertices = []
|
|
66
|
+
faces = []
|
|
67
|
+
for line in content.splitlines():
|
|
68
|
+
if line.startswith("v "):
|
|
69
|
+
parts = line.split()
|
|
70
|
+
vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
|
|
71
|
+
elif line.startswith("f "):
|
|
72
|
+
parts = line.split()
|
|
73
|
+
faces.append([int(parts[1])-1, int(parts[2])-1, int(parts[3])-1])
|
|
74
|
+
|
|
75
|
+
vertices_t = torch.tensor(vertices, dtype=torch.float32)
|
|
76
|
+
faces_t = torch.tensor(faces, dtype=torch.int32)
|
|
77
|
+
|
|
78
|
+
if self.cached:
|
|
79
|
+
self._cache[idx] = (vertices_t, faces_t, label)
|
|
80
|
+
|
|
81
|
+
if self.transform:
|
|
82
|
+
vertices_t, faces_t = self.transform(vertices_t.clone(), faces_t.clone())
|
|
83
|
+
|
|
84
|
+
return vertices_t, faces_t, label
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class SparseDigit3D(Digit3D):
|
|
88
|
+
"""
|
|
89
|
+
Digit3D Dataset that constructs a sparse SDF voxel grid from the mesh on-the-fly.
|
|
90
|
+
"""
|
|
91
|
+
def __init__(self, root: str = "~/.conquer3d/", train: bool = True, transform=None, download: bool = False,
|
|
92
|
+
grid_res: int = 32, grid_bound: float = 1.2, cached: bool = False):
|
|
93
|
+
super().__init__(root, train, transform, download, cached=cached)
|
|
94
|
+
self.grid_res = grid_res
|
|
95
|
+
self.grid_bound = grid_bound
|
|
96
|
+
|
|
97
|
+
def __getitem__(self, idx: int):
|
|
98
|
+
# 1. Obtain vertices, faces, and label from Digit3D
|
|
99
|
+
vertices, faces, label = super().__getitem__(idx)
|
|
100
|
+
|
|
101
|
+
# 2. Construct voxel grid in CPU
|
|
102
|
+
grid_vertices, voxels, idx_grids = c3d.data_structure.create_voxel_grid(
|
|
103
|
+
grid_min=[-self.grid_bound] * 3,
|
|
104
|
+
grid_max=[self.grid_bound] * 3,
|
|
105
|
+
res=[self.grid_res] * 3,
|
|
106
|
+
device="cpu"
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
# Construct the mesh using meshlib's numpy interface
|
|
110
|
+
mesh_mr = mrnp.meshFromFacesVerts(faces.numpy(), vertices.numpy())
|
|
111
|
+
|
|
112
|
+
# Construct point cloud from grid vertices for vectorized distance computation
|
|
113
|
+
pc = mrnp.pointCloudFromPoints(grid_vertices.numpy())
|
|
114
|
+
|
|
115
|
+
# Compute the signed distance for all grid vertices at once
|
|
116
|
+
dist_scalars = mr.findSignedDistances(mesh_mr, pc.points)
|
|
117
|
+
|
|
118
|
+
# Convert to tensor
|
|
119
|
+
sdf = torch.tensor(list(dist_scalars), dtype=torch.float32, device="cpu")
|
|
120
|
+
|
|
121
|
+
# 4. Compute active voxels
|
|
122
|
+
active_voxel_indices = c3d.data_structure.compute_active_voxels(voxels, sdf, iso=0.0)
|
|
123
|
+
|
|
124
|
+
# 5. Extract purely Voxel-Centric representations using voxel2sparse
|
|
125
|
+
sparse_coords, sparse_sdfs = c3d.conversion.grid.voxel2sparse(
|
|
126
|
+
active_voxel_indices, voxels, idx_grids, sdf=sdf, batch_idx=0
|
|
127
|
+
)
|
|
128
|
+
|
|
129
|
+
# We only need the x, y, z for the dataset (collate_fn handles batch_idx)
|
|
130
|
+
sparse_idx_grids = sparse_coords[:, 1:]
|
|
131
|
+
|
|
132
|
+
return sparse_idx_grids, sparse_sdfs, label
|
|
File without changes
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
import random
|
|
2
|
+
|
|
3
|
+
class BaseTransform(object):
|
|
4
|
+
def __init__(self, p=1.0):
|
|
5
|
+
self.p = p
|
|
6
|
+
|
|
7
|
+
def __call__(self, **data):
|
|
8
|
+
if random.random() <= self.p:
|
|
9
|
+
return self.apply(**data)
|
|
10
|
+
return data
|
|
11
|
+
|
|
12
|
+
def apply(self, **data):
|
|
13
|
+
raise NotImplementedError("Transform must implement apply")
|
|
14
|
+
|
|
15
|
+
class Sequence(BaseTransform):
|
|
16
|
+
def __init__(self, transforms, shuffle=False, p=1.0):
|
|
17
|
+
super().__init__(p=p)
|
|
18
|
+
self.transforms = transforms
|
|
19
|
+
self.shuffle = shuffle
|
|
20
|
+
|
|
21
|
+
def apply(self, **data):
|
|
22
|
+
transforms_to_apply = self.transforms.copy()
|
|
23
|
+
if self.shuffle:
|
|
24
|
+
random.shuffle(transforms_to_apply)
|
|
25
|
+
|
|
26
|
+
for transform in transforms_to_apply:
|
|
27
|
+
data = transform(**data)
|
|
28
|
+
|
|
29
|
+
return data
|
|
30
|
+
|
|
31
|
+
class MeshSequence(Sequence):
|
|
32
|
+
def __call__(self, vertices, faces):
|
|
33
|
+
if random.random() <= self.p:
|
|
34
|
+
return self.apply(vertices, faces)
|
|
35
|
+
return vertices, faces
|
|
36
|
+
|
|
37
|
+
def apply(self, vertices, faces):
|
|
38
|
+
data = super().apply(vertices=vertices, faces=faces)
|
|
39
|
+
return data.get('vertices', vertices), data.get('faces', faces)
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import math
|
|
3
|
+
|
|
4
|
+
def rotation(vertices, rotation_axis, rotation_degree):
|
|
5
|
+
"""
|
|
6
|
+
Rotates vertices around a given axis by a specific degree.
|
|
7
|
+
|
|
8
|
+
Args:
|
|
9
|
+
vertices (torch.Tensor): Shape (N, 3) representing 3D coordinates.
|
|
10
|
+
rotation_axis (str): 'x', 'y', or 'z'.
|
|
11
|
+
rotation_degree (float): Degree to rotate by (in degrees).
|
|
12
|
+
|
|
13
|
+
Returns:
|
|
14
|
+
torch.Tensor: The rotated vertices of the same type.
|
|
15
|
+
"""
|
|
16
|
+
theta = math.radians(rotation_degree)
|
|
17
|
+
cos_t = math.cos(theta)
|
|
18
|
+
sin_t = math.sin(theta)
|
|
19
|
+
|
|
20
|
+
device = vertices.device
|
|
21
|
+
dtype = vertices.dtype
|
|
22
|
+
|
|
23
|
+
# Define rotation matrices based on the axis directly as PyTorch tensors
|
|
24
|
+
if rotation_axis == 'x':
|
|
25
|
+
rot_mat = torch.tensor([
|
|
26
|
+
[1.0, 0.0, 0.0],
|
|
27
|
+
[0.0, cos_t, -sin_t],
|
|
28
|
+
[0.0, sin_t, cos_t]
|
|
29
|
+
], dtype=dtype, device=device)
|
|
30
|
+
elif rotation_axis == 'y':
|
|
31
|
+
rot_mat = torch.tensor([
|
|
32
|
+
[cos_t, 0.0, sin_t],
|
|
33
|
+
[0.0, 1.0, 0.0],
|
|
34
|
+
[-sin_t, 0.0, cos_t]
|
|
35
|
+
], dtype=dtype, device=device)
|
|
36
|
+
elif rotation_axis == 'z':
|
|
37
|
+
rot_mat = torch.tensor([
|
|
38
|
+
[cos_t, -sin_t, 0.0],
|
|
39
|
+
[sin_t, cos_t, 0.0],
|
|
40
|
+
[0.0, 0.0, 1.0]
|
|
41
|
+
], dtype=dtype, device=device)
|
|
42
|
+
else:
|
|
43
|
+
raise ValueError("rotation_axis must be 'x', 'y', or 'z'")
|
|
44
|
+
|
|
45
|
+
# Compute dot product
|
|
46
|
+
return torch.matmul(vertices, rot_mat.T)
|
|
47
|
+
|
|
48
|
+
def scale(vertices, scale_factor):
|
|
49
|
+
"""
|
|
50
|
+
Scales vertices by a given factor.
|
|
51
|
+
|
|
52
|
+
Args:
|
|
53
|
+
vertices (torch.Tensor): Shape (N, 3) representing 3D coordinates.
|
|
54
|
+
scale_factor (float, tuple, or list): The factor(s) to scale by.
|
|
55
|
+
|
|
56
|
+
Returns:
|
|
57
|
+
torch.Tensor: The scaled vertices.
|
|
58
|
+
"""
|
|
59
|
+
if isinstance(scale_factor, (list, tuple)):
|
|
60
|
+
scale_factor = torch.tensor(scale_factor, dtype=vertices.dtype, device=vertices.device)
|
|
61
|
+
|
|
62
|
+
return vertices * scale_factor
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import random
|
|
3
|
+
from .base import BaseTransform
|
|
4
|
+
from .ops import rotation, scale
|
|
5
|
+
|
|
6
|
+
class Rotation(BaseTransform):
|
|
7
|
+
def __init__(self, axis='z', degree_range=(-180, 180), p=1.0):
|
|
8
|
+
"""
|
|
9
|
+
Args:
|
|
10
|
+
axis (str or list): The axis or list of axes to rotate around ('x', 'y', 'z').
|
|
11
|
+
degree_range (tuple): Range of degrees to randomly sample from (min, max).
|
|
12
|
+
p (float): Probability of applying the transform.
|
|
13
|
+
"""
|
|
14
|
+
super().__init__(p=p)
|
|
15
|
+
self.axis = axis
|
|
16
|
+
self.degree_range = degree_range
|
|
17
|
+
|
|
18
|
+
def apply(self, **data):
|
|
19
|
+
if 'vertices' not in data:
|
|
20
|
+
raise KeyError("Rotation transform requires 'vertices' in data")
|
|
21
|
+
|
|
22
|
+
vertices = data['vertices']
|
|
23
|
+
|
|
24
|
+
# Support multiple axes in sequence
|
|
25
|
+
axes = self.axis if isinstance(self.axis, (list, tuple)) else [self.axis]
|
|
26
|
+
|
|
27
|
+
for ax in axes:
|
|
28
|
+
# Sample a random degree for each axis
|
|
29
|
+
degree = random.uniform(self.degree_range[0], self.degree_range[1])
|
|
30
|
+
vertices = rotation(vertices, ax, degree)
|
|
31
|
+
|
|
32
|
+
data['vertices'] = vertices
|
|
33
|
+
return data
|
|
34
|
+
|
|
35
|
+
class Scale(BaseTransform):
|
|
36
|
+
def __init__(self, scale_range=(0.8, 1.2), p=1.0):
|
|
37
|
+
"""
|
|
38
|
+
Args:
|
|
39
|
+
scale_range (tuple): Range of scale factors to randomly sample from (min, max).
|
|
40
|
+
p (float): Probability of applying the transform.
|
|
41
|
+
"""
|
|
42
|
+
super().__init__(p=p)
|
|
43
|
+
self.scale_range = scale_range
|
|
44
|
+
|
|
45
|
+
def apply(self, **data):
|
|
46
|
+
if 'vertices' not in data:
|
|
47
|
+
raise KeyError("Scale transform requires 'vertices' in data")
|
|
48
|
+
|
|
49
|
+
vertices = data['vertices']
|
|
50
|
+
|
|
51
|
+
# Sample a uniform random scale factor
|
|
52
|
+
scale_factor = random.uniform(self.scale_range[0], self.scale_range[1])
|
|
53
|
+
|
|
54
|
+
data['vertices'] = scale(vertices, scale_factor)
|
|
55
|
+
return data
|
|
@@ -6,10 +6,18 @@ from .._C import (
|
|
|
6
6
|
MeshBVH,
|
|
7
7
|
TriangleMesh
|
|
8
8
|
)
|
|
9
|
-
from .grid import
|
|
9
|
+
from .grid import (
|
|
10
|
+
create_voxel_grid,
|
|
11
|
+
compute_grid_normal,
|
|
12
|
+
compute_active_voxels,
|
|
13
|
+
)
|
|
10
14
|
|
|
11
15
|
spatial_data_structures = ['KDTree', 'BVH', 'GSBVH', 'PGSBVH', 'MeshBVH']
|
|
12
16
|
mesh_data_structures = ['TriangleMesh']
|
|
13
|
-
grid_data_structures = [
|
|
17
|
+
grid_data_structures = [
|
|
18
|
+
'create_voxel_grid',
|
|
19
|
+
'compute_grid_normal',
|
|
20
|
+
'compute_active_voxels',
|
|
21
|
+
]
|
|
14
22
|
|
|
15
23
|
__all__ = spatial_data_structures + mesh_data_structures + grid_data_structures
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
|
|
3
|
+
class BTriangleMesh:
|
|
4
|
+
def __init__(self, vertices: torch.Tensor, faces: torch.Tensor, vertbids: torch.Tensor, facebids: torch.Tensor, batch_size: int):
|
|
5
|
+
self.vertices = vertices
|
|
6
|
+
self.faces = faces
|
|
7
|
+
self.vertbids = vertbids
|
|
8
|
+
self.facebids = facebids
|
|
9
|
+
self.batch_size = batch_size
|
|
10
|
+
|
|
11
|
+
def to(self, device):
|
|
12
|
+
self.vertices = self.vertices.to(device)
|
|
13
|
+
self.faces = self.faces.to(device)
|
|
14
|
+
self.vertbids = self.vertbids.to(device)
|
|
15
|
+
self.facebids = self.facebids.to(device)
|
|
16
|
+
return self
|
|
17
|
+
|
|
18
|
+
def cuda(self, non_blocking=False):
|
|
19
|
+
self.vertices = self.vertices.cuda(non_blocking=non_blocking)
|
|
20
|
+
self.faces = self.faces.cuda(non_blocking=non_blocking)
|
|
21
|
+
self.vertbids = self.vertbids.cuda(non_blocking=non_blocking)
|
|
22
|
+
self.facebids = self.facebids.cuda(non_blocking=non_blocking)
|
|
23
|
+
return self
|
|
@@ -62,53 +62,4 @@ def compute_active_voxels(
|
|
|
62
62
|
Returns:
|
|
63
63
|
torch.Tensor: An int64 tensor containing the indices of the active voxels.
|
|
64
64
|
"""
|
|
65
|
-
return _C.compute_active_voxels(voxels.contiguous(), sdf.contiguous(), iso)
|
|
66
|
-
|
|
67
|
-
def voxel2sparse(
|
|
68
|
-
active_indices: torch.Tensor,
|
|
69
|
-
voxels: torch.Tensor,
|
|
70
|
-
idx_grids: torch.Tensor,
|
|
71
|
-
batch_idx: int = 0
|
|
72
|
-
) -> torch.Tensor:
|
|
73
|
-
"""
|
|
74
|
-
Converts 1D active voxel indices into the [N, 4] coordinates tensor
|
|
75
|
-
expected by sparse neural networks (like spconv/torchsparse).
|
|
76
|
-
|
|
77
|
-
Args:
|
|
78
|
-
active_indices (torch.Tensor): 1D tensor of active voxel indices.
|
|
79
|
-
voxels (torch.Tensor): An int32 tensor of shape (V, 8) containing corner indices for each voxel.
|
|
80
|
-
idx_grids (torch.Tensor): An int64 tensor of shape (N, 3) containing the (i, j, k) 3D coordinate indices for each vertex.
|
|
81
|
-
batch_idx (int, optional): The batch index to prepend. Defaults to 0.
|
|
82
|
-
|
|
83
|
-
Returns:
|
|
84
|
-
torch.Tensor: An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
|
|
85
|
-
"""
|
|
86
|
-
active_corners = voxels[active_indices]
|
|
87
|
-
base_vertex_indices = active_corners[:, 0].to(torch.int64)
|
|
88
|
-
voxel_coords = idx_grids[base_vertex_indices]
|
|
89
|
-
|
|
90
|
-
batch_col = torch.full((voxel_coords.shape[0], 1), batch_idx, device=voxel_coords.device, dtype=torch.int64)
|
|
91
|
-
sparse_coords = torch.cat([batch_col, voxel_coords], dim=1).to(torch.int32)
|
|
92
|
-
return sparse_coords
|
|
93
|
-
|
|
94
|
-
def sparse2voxel(
|
|
95
|
-
sparse_coords: torch.Tensor,
|
|
96
|
-
res: Union[List[int], Tuple[int, int, int]]
|
|
97
|
-
) -> torch.Tensor:
|
|
98
|
-
"""
|
|
99
|
-
Converts a [N, 4] sparse coordinates tensor back to 1D active voxel indices.
|
|
100
|
-
|
|
101
|
-
Args:
|
|
102
|
-
sparse_coords (torch.Tensor): An int32 tensor of shape (N, 4) in the format [batch_idx, x, y, z].
|
|
103
|
-
res (List[int] | Tuple[int, int, int]): The number of vertices along each axis (rx, ry, rz).
|
|
104
|
-
|
|
105
|
-
Returns:
|
|
106
|
-
torch.Tensor: An int64 tensor containing the indices of the active voxels.
|
|
107
|
-
"""
|
|
108
|
-
x = sparse_coords[:, 1].to(torch.int64)
|
|
109
|
-
y = sparse_coords[:, 2].to(torch.int64)
|
|
110
|
-
z = sparse_coords[:, 3].to(torch.int64)
|
|
111
|
-
|
|
112
|
-
rx, ry, rz = res
|
|
113
|
-
active_indices = x * (ry - 1) * (rz - 1) + y * (rz - 1) + z
|
|
114
|
-
return active_indices
|
|
65
|
+
return _C.compute_active_voxels(voxels.contiguous(), sdf.contiguous(), iso)
|
|
File without changes
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: conquer3d
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: Geometric Cuda Tool Box
|
|
5
5
|
Author-email: Do Hoang Khoi <khoido8899@gmail.com>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -13,6 +13,7 @@ Requires-Python: >=3.9
|
|
|
13
13
|
Description-Content-Type: text/markdown
|
|
14
14
|
License-File: LICENSE
|
|
15
15
|
Requires-Dist: torch>=2.0.0
|
|
16
|
+
Requires-Dist: gdown
|
|
16
17
|
Provides-Extra: cuda
|
|
17
18
|
Requires-Dist: cuda-toolkit==12.8.2; extra == "cuda"
|
|
18
19
|
Provides-Extra: dev
|
|
@@ -61,7 +62,9 @@ pip install kaolin==0.18.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/t
|
|
|
61
62
|
pip install pybind11-stubgen
|
|
62
63
|
pip install git+https://github.com/KhoiDOO/geocutool.git --no-build-isolation
|
|
63
64
|
|
|
64
|
-
pip install plotly open3d jupyter trimesh point-cloud-utils
|
|
65
|
+
pip install plotly open3d jupyter trimesh point-cloud-utils meshlib
|
|
66
|
+
|
|
67
|
+
pip install rectified-flow-pytorch
|
|
65
68
|
```
|
|
66
69
|
|
|
67
70
|
# Development
|
|
@@ -12,6 +12,9 @@ conquer3d.egg-info/dependency_links.txt
|
|
|
12
12
|
conquer3d.egg-info/not-zip-safe
|
|
13
13
|
conquer3d.egg-info/requires.txt
|
|
14
14
|
conquer3d.egg-info/top_level.txt
|
|
15
|
+
conquer3d/conversion/__init__.py
|
|
16
|
+
conquer3d/conversion/grid.py
|
|
17
|
+
conquer3d/conversion/mesh.py
|
|
15
18
|
conquer3d/creation/__init__.py
|
|
16
19
|
conquer3d/creation/triangle_creation.py
|
|
17
20
|
conquer3d/csrc/check.h
|
|
@@ -66,7 +69,19 @@ conquer3d/csrc/primitive/pgs_aabb.cu
|
|
|
66
69
|
conquer3d/csrc/primitive/pgs_math.cuh
|
|
67
70
|
conquer3d/csrc/primitive/ray.h
|
|
68
71
|
conquer3d/csrc/primitive/triangles.h
|
|
72
|
+
conquer3d/data/__init__.py
|
|
73
|
+
conquer3d/data/collate/__init__.py
|
|
74
|
+
conquer3d/data/collate/mesh.py
|
|
75
|
+
conquer3d/data/collate/sparse_tensor.py
|
|
76
|
+
conquer3d/data/dataset/__init__.py
|
|
77
|
+
conquer3d/data/dataset/base_mesh.py
|
|
78
|
+
conquer3d/data/dataset/digit3d.py
|
|
79
|
+
conquer3d/data/transform/__init__.py
|
|
80
|
+
conquer3d/data/transform/base.py
|
|
81
|
+
conquer3d/data/transform/ops.py
|
|
82
|
+
conquer3d/data/transform/vertex.py
|
|
69
83
|
conquer3d/data_structure/__init__.py
|
|
84
|
+
conquer3d/data_structure/bmesh.py
|
|
70
85
|
conquer3d/data_structure/grid.py
|
|
71
86
|
conquer3d/ops/__init__.py
|
|
72
87
|
conquer3d/ops/diff_marching_cubes.py
|
|
@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
|
|
|
9
9
|
|
|
10
10
|
[project]
|
|
11
11
|
name = "conquer3d"
|
|
12
|
-
version = "0.3.
|
|
12
|
+
version = "0.3.2"
|
|
13
13
|
description = "Geometric Cuda Tool Box"
|
|
14
14
|
readme = "README.md"
|
|
15
15
|
license = "MIT"
|
|
@@ -20,6 +20,7 @@ requires-python = ">=3.9"
|
|
|
20
20
|
keywords = ["geometry"]
|
|
21
21
|
dependencies = [
|
|
22
22
|
"torch>=2.0.0",
|
|
23
|
+
"gdown",
|
|
23
24
|
]
|
|
24
25
|
|
|
25
26
|
[project.optional-dependencies]
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|