423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-12-19 05:40:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python marching cubes improvements

Browse Source 
Summary: Overhaul of marching_cubes_naive for better performance and to avoid relying on unstable hashing. In particular, instead of hashing vertex positions, we index each interpolated vertex with its corresponding edge in the 3d grid.

Reviewed By: kjchalup

Differential Revision: D39419642

fbshipit-source-id: b5fede3525c545d1d374198928dfb216262f0ec0
This commit is contained in:
Jiali Duan
2022-10-06 11:08:49 -07:00
committed by Facebook GitHub Bot
parent 6471893f59
commit 850efdf706
5 changed files with 342 additions and 703 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
tests/data/test_marching_cubes_data/double_ellipsoid.pickle
View File

Binary file not shown.

BIN
tests/data/test_marching_cubes_data/sphere_level64.pickle
View File

Binary file not shown.

321
tests/test_marching_cubes.py
View File

@@ -32,8 +32,8 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=False)
expected_verts = torch.tensor([])
expected_faces = torch.tensor([], dtype=torch.int64)
expected_verts = torch.tensor([[]])
expected_faces = torch.tensor([[]], dtype=torch.int64)
self.assertClose(verts, expected_verts)
self.assertClose(faces, expected_faces)
@@ -42,16 +42,15 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
volume_data[0, 0, 0, 0] = 0
volume_data = volume_data.permute(0, 3, 2, 1) # (B, D, H, W)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=False)
expected_verts = torch.tensor(
[
[0.5, 0, 0],
[0, 0, 0.5],
[0, 0.5, 0],
[0, 0, 0.5],
]
)
expected_faces = torch.tensor([[1, 2, 0]])
expected_faces = torch.tensor([[0, 1, 2]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -69,12 +68,12 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[1.0000, 0.0000, 0.5000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
[1.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 0.0000],
]
)
expected_faces = torch.tensor([[1, 2, 0], [3, 2, 1]])
expected_faces = torch.tensor([[0, 1, 2], [3, 1, 0]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -92,15 +91,15 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[1.0000, 0.5000, 0.0000],
[1.0000, 1.0000, 0.5000],
[0.5000, 1.0000, 0.0000],
[1.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
]
)
expected_faces = torch.tensor([[0, 1, 5], [4, 3, 2]])
expected_faces = torch.tensor([[0, 1, 2], [3, 4, 5]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -119,14 +118,14 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[1.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
]
)
expected_faces = torch.tensor([[0, 2, 1], [0, 4, 2], [4, 3, 2]])
expected_faces = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 4, 1]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -143,15 +142,14 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 0.0000],
[1.0000, 0.5000, 1.0000],
[0.0000, 0.5000, 1.0000],
]
)
expected_faces = torch.tensor([[1, 0, 2], [2, 0, 3]])
expected_faces = torch.tensor([[0, 1, 2], [2, 1, 3]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -171,17 +169,17 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 0.0000],
[0.0000, 1.0000, 0.5000],
[0.0000, 0.5000, 0.0000],
[1.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
]
)
expected_faces = torch.tensor([[2, 7, 3], [0, 6, 1], [6, 4, 1], [6, 5, 4]])
expected_faces = torch.tensor([[0, 1, 2], [3, 4, 5], [3, 5, 6], [5, 4, 7]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -202,22 +200,22 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 1.0000],
[1.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[1.0000, 1.0000, 0.5000],
[0.5000, 1.0000, 0.0000],
[0.0000, 1.0000, 0.5000],
[0.0000, 0.5000, 1.0000],
[0.0000, 1.0000, 0.5000],
[1.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 1.0000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 0.0000],
[1.0000, 0.5000, 0.0000],
[1.0000, 1.0000, 0.5000],
]
)
expected_faces = torch.tensor([[0, 1, 9], [4, 7, 8], [2, 3, 11], [5, 10, 6]])
expected_faces = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -238,15 +236,15 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[1.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 0.0000],
[0.5000, 1.0000, 1.0000],
[0.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[0.0000, 1.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[1.0000, 0.0000, 0.5000],
[0.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
]
)
expected_faces = torch.tensor([[2, 3, 5], [4, 2, 5], [4, 5, 1], [4, 1, 0]])
expected_faces = torch.tensor([[0, 1, 2], [3, 1, 0], [3, 4, 1], [3, 5, 4]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -269,13 +267,13 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
[
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[0.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.5000, 1.0000, 1.0000],
]
)
expected_faces = torch.tensor([[0, 5, 4], [0, 4, 3], [0, 3, 1], [3, 4, 2]])
expected_faces = torch.tensor([[0, 1, 2], [0, 2, 3], [0, 3, 4], [5, 3, 2]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -295,15 +293,15 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 0.0000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[1.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[0.0000, 0.5000, 0.0000],
[0.5000, 1.0000, 1.0000],
]
)
expected_faces = torch.tensor([[4, 3, 2], [0, 1, 5]])
expected_faces = torch.tensor([[0, 1, 2], [3, 4, 5]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -324,16 +322,16 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[1.0000, 0.0000, 0.5000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[1.0000, 0.5000, 0.0000],
[1.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 0.0000],
[0.5000, 1.0000, 1.0000],
]
)
expected_faces = torch.tensor([[5, 1, 6], [5, 0, 1], [4, 3, 2]])
expected_faces = torch.tensor([[0, 1, 2], [0, 3, 1], [4, 5, 6]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -354,18 +352,18 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[1.0000, 0.0000, 0.5000],
[1.0000, 0.5000, 0.0000],
[0.5000, 0.0000, 0.0000],
[0.5000, 1.0000, 1.0000],
[1.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[0.5000, 1.0000, 1.0000],
[0.0000, 0.5000, 0.0000],
[0.5000, 1.0000, 0.0000],
[0.0000, 1.0000, 0.5000],
[1.0000, 0.5000, 1.0000],
[1.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 0.0000],
]
)
expected_faces = torch.tensor([[6, 3, 2], [7, 0, 1], [5, 4, 8]])
expected_faces = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -386,18 +384,18 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5000, 0.0000, 1.0000],
[1.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[0.5000, 0.0000, 1.0000],
[1.0000, 1.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[0.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 0.0000],
[0.5000, 1.0000, 0.0000],
[0.0000, 1.0000, 0.5000],
]
)
expected_faces = torch.tensor([[3, 6, 2], [3, 7, 6], [1, 5, 0], [5, 4, 0]])
expected_faces = torch.tensor([[0, 1, 2], [2, 1, 3], [4, 5, 6], [4, 6, 7]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -418,16 +416,16 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[1.0000, 0.0000, 0.5000],
[0.5000, 0.0000, 0.0000],
[0.5000, 1.0000, 1.0000],
[1.0000, 1.0000, 0.5000],
[0.0000, 0.5000, 1.0000],
[0.0000, 0.5000, 0.0000],
[0.0000, 0.5000, 1.0000],
[1.0000, 1.0000, 0.5000],
[1.0000, 0.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
]
)
expected_faces = torch.tensor([[1, 0, 3], [1, 3, 4], [1, 4, 5], [2, 4, 3]])
expected_faces = torch.tensor([[0, 1, 2], [0, 2, 3], [0, 3, 4], [3, 2, 5]])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -447,27 +445,26 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.5, 1, 1],
[1, 1, 0.5],
[1, 0.5, 1],
[1, 1, 1.5],
[1, 1.5, 1],
[1.5, 1, 1],
[1.0000, 0.5000, 1.0000],
[1.0000, 1.0000, 0.5000],
[0.5000, 1.0000, 1.0000],
[1.5000, 1.0000, 1.0000],
[1.0000, 1.5000, 1.0000],
[1.0000, 1.0000, 1.5000],
]
)
expected_faces = torch.tensor(
[
[2, 0, 1],
[2, 3, 0],
[0, 4, 1],
[3, 4, 0],
[5, 2, 1],
[3, 2, 5],
[5, 1, 4],
[0, 1, 2],
[1, 0, 3],
[1, 4, 2],
[1, 3, 4],
[0, 2, 5],
[3, 0, 5],
[2, 4, 5],
[3, 5, 4],
]
)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -492,76 +489,76 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[0.9000, 1.0000, 1.0000],
[1.0000, 1.0000, 0.9000],
[1.0000, 0.9000, 1.0000],
[0.9000, 1.0000, 2.0000],
[1.0000, 0.9000, 2.0000],
[1.0000, 1.0000, 2.1000],
[0.9000, 2.0000, 1.0000],
[1.0000, 2.0000, 0.9000],
[0.9000, 2.0000, 2.0000],
[1.0000, 2.0000, 2.1000],
[1.0000, 2.1000, 1.0000],
[1.0000, 2.1000, 2.0000],
[2.0000, 1.0000, 0.9000],
[1.0000, 1.0000, 0.9000],
[0.9000, 1.0000, 1.0000],
[2.0000, 0.9000, 1.0000],
[2.0000, 0.9000, 2.0000],
[2.0000, 1.0000, 2.1000],
[2.0000, 2.0000, 0.9000],
[2.0000, 2.0000, 2.1000],
[2.0000, 2.1000, 1.0000],
[2.0000, 2.1000, 2.0000],
[2.0000, 1.0000, 0.9000],
[2.1000, 1.0000, 1.0000],
[2.1000, 1.0000, 2.0000],
[1.0000, 2.0000, 0.9000],
[0.9000, 2.0000, 1.0000],
[2.0000, 2.0000, 0.9000],
[2.1000, 2.0000, 1.0000],
[1.0000, 2.1000, 1.0000],
[2.0000, 2.1000, 1.0000],
[1.0000, 0.9000, 2.0000],
[0.9000, 1.0000, 2.0000],
[2.0000, 0.9000, 2.0000],
[2.1000, 1.0000, 2.0000],
[0.9000, 2.0000, 2.0000],
[2.1000, 2.0000, 2.0000],
[1.0000, 2.1000, 2.0000],
[2.0000, 2.1000, 2.0000],
[1.0000, 1.0000, 2.1000],
[2.0000, 1.0000, 2.1000],
[1.0000, 2.0000, 2.1000],
[2.0000, 2.0000, 2.1000],
]
)
expected_faces = torch.tensor(
[
[2, 0, 1],
[2, 4, 3],
[0, 2, 3],
[4, 5, 3],
[0, 6, 7],
[1, 0, 7],
[3, 8, 0],
[8, 6, 0],
[5, 9, 8],
[3, 5, 8],
[0, 1, 2],
[0, 3, 4],
[1, 0, 4],
[4, 3, 5],
[1, 6, 7],
[2, 1, 7],
[4, 8, 1],
[1, 8, 6],
[8, 4, 5],
[9, 8, 5],
[6, 10, 7],
[11, 10, 6],
[8, 11, 6],
[9, 11, 8],
[13, 2, 1],
[12, 13, 1],
[14, 4, 13],
[13, 4, 2],
[4, 14, 15],
[5, 4, 15],
[12, 1, 16],
[1, 7, 16],
[15, 17, 5],
[5, 17, 9],
[16, 7, 10],
[18, 16, 10],
[19, 18, 11],
[18, 10, 11],
[6, 8, 11],
[10, 6, 11],
[8, 9, 11],
[12, 0, 2],
[13, 12, 2],
[3, 0, 14],
[14, 0, 12],
[15, 5, 3],
[14, 15, 3],
[2, 7, 13],
[7, 16, 13],
[5, 15, 9],
[9, 15, 17],
[10, 18, 16],
[7, 10, 16],
[11, 19, 10],
[19, 18, 10],
[9, 17, 19],
[11, 9, 19],
[20, 13, 12],
[20, 21, 14],
[13, 20, 14],
[12, 13, 20],
[14, 12, 20],
[21, 14, 20],
[15, 14, 21],
[22, 20, 12],
[16, 22, 12],
[13, 16, 22],
[20, 13, 22],
[21, 20, 23],
[23, 20, 22],
[20, 22, 23],
[17, 15, 21],
[23, 17, 21],
[22, 16, 18],
[16, 18, 22],
[23, 22, 18],
[19, 23, 18],
[17, 23, 19],
@@ -569,6 +566,7 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, 0.9, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 5)
self.assertClose(verts[0], expected_verts)
@@ -592,34 +590,49 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
expected_verts = torch.tensor(
[
[2.0, 1.0, 1.0],
[2.0, 2.0, 1.0],
[1.0, 1.0, 1.0],
[1.0, 1.0, 2.0],
[1.0, 2.0, 1.0],
[2.0, 1.0, 1.0],
[1.0, 1.0, 1.0],
[2.0, 1.0, 2.0],
[1.0, 1.0, 2.0],
[1.0, 1.0, 1.0],
[1.0, 2.0, 1.0],
[1.0, 1.0, 2.0],
[1.0, 2.0, 2.0],
[2.0, 1.0, 1.0],
[2.0, 1.0, 2.0],
[2.0, 2.0, 1.0],
[2.0, 2.0, 2.0],
[2.0, 2.0, 1.0],
[2.0, 2.0, 2.0],
[1.0, 2.0, 1.0],
[1.0, 2.0, 2.0],
[2.0, 1.0, 2.0],
[1.0, 1.0, 2.0],
[2.0, 2.0, 2.0],
[1.0, 2.0, 2.0],
]
)
expected_faces = torch.tensor(
[
[1, 3, 0],
[3, 2, 0],
[5, 1, 4],
[4, 1, 0],
[4, 0, 6],
[0, 2, 6],
[5, 7, 1],
[1, 7, 3],
[7, 6, 3],
[6, 2, 3],
[5, 4, 7],
[7, 4, 6],
[0, 1, 2],
[2, 1, 3],
[4, 5, 6],
[6, 5, 7],
[8, 9, 10],
[9, 11, 10],
[12, 13, 14],
[14, 13, 15],
[16, 17, 18],
[17, 19, 18],
[20, 21, 22],
[21, 23, 22],
]
)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -651,8 +664,8 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
filename = os.path.join(DATA_DIR, data_filename)
with open(filename, "rb") as file:
verts_and_faces = pickle.load(file)
expected_verts = verts_and_faces["verts"].squeeze()
expected_faces = verts_and_faces["faces"].squeeze()
expected_verts = verts_and_faces["verts"]
expected_faces = verts_and_faces["faces"]
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
@@ -689,8 +702,8 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
expected_verts = verts_and_faces["verts"]
expected_faces = verts_and_faces["faces"]
self.assertClose(verts[0], expected_verts[0])
self.assertClose(faces[0], expected_faces[0])
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_cube_surface_area(self):
if USE_SCIKIT:
@@ -760,16 +773,26 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
self.assertClose(surf, surf_sci)
def test_ball_example(self):
N = 15
axis_tensor = torch.arange(0, N)
X, Y, Z = torch.meshgrid(axis_tensor, axis_tensor, axis_tensor, indexing="ij")
u = (X - 15) ** 2 + (Y - 15) ** 2 + (Z - 15) ** 2 - 8**2
u = u[None].float()
verts, faces = marching_cubes_naive(u, 0, return_local_coords=False)
@staticmethod
def marching_cubes_with_init(batch_size: int, V: int):
def marching_cubes_with_init(algo_type: str, batch_size: int, V: int):
device = torch.device("cuda:0")
volume_data = torch.rand(
(batch_size, V, V, V), dtype=torch.float32, device=device
)
torch.cuda.synchronize()
algo_table = {
"naive": marching_cubes_naive,
}
def convert():
marching_cubes_naive(volume_data, return_local_coords=False)
algo_table[algo_type](volume_data, return_local_coords=False)
torch.cuda.synchronize()
return convert