423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2026-01-17 11:50:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Marching Cubes C++ torch extension

Browse Source 
Summary:
Torch C++ extension for Marching Cubes

- Add torch C++ extension for marching cubes. Observe a speed up of ~255x-324x speed up (over varying batch sizes and spatial resolutions)

- Add C++ impl in existing unit-tests.

(Note: this ignores all push blocking failures!)

Reviewed By: kjchalup

Differential Revision: D39590638

fbshipit-source-id: e44d2852a24c2c398e5ea9db20f0dfaa1817e457
This commit is contained in:
Jiali Duan
2022-10-06 11:13:53 -07:00
committed by Facebook GitHub Bot
parent 850efdf706
commit 0d8608b9f9
7 changed files with 879 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

196
tests/test_marching_cubes.py
View File

@@ -9,7 +9,7 @@ import pickle
import unittest
import torch
from pytorch3d.ops.marching_cubes import marching_cubes_naive
from pytorch3d.ops.marching_cubes import marching_cubes, marching_cubes_naive
from .common_testing import get_tests_dir, TestCaseMixin
@@ -37,6 +37,11 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts, expected_verts)
self.assertClose(faces, expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts, expected_verts)
self.assertClose(faces, expected_faces)
def test_case1(self): # case 1
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -54,11 +59,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case2(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0:2, 0, 0] = 0
@@ -77,11 +92,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case3(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -103,11 +128,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case4(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 1, 0, 0] = 0
@@ -129,11 +164,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case5(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0:2, 0, 0:2] = 0
@@ -153,11 +198,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case6(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 1, 0, 0] = 0
@@ -184,11 +239,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case7(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -220,11 +285,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case8(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -249,11 +324,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case9(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 1, 0, 0] = 0
@@ -278,11 +363,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case10(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -306,11 +401,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case11(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -336,11 +441,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case12(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 1, 0, 0] = 0
@@ -368,11 +483,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case13(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -400,11 +525,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
def test_case14(self):
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
volume_data[0, 0, 0, 0] = 0
@@ -430,11 +565,21 @@ class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 2)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
def test_single_point(self):
@@ -468,12 +613,23 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=True)
expected_verts = convert_to_local(expected_verts, 3)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
# test C++ implementation
verts, faces = marching_cubes(volume_data, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
def test_cube(self):
volume_data = torch.zeros(1, 5, 5, 5) # (B, W, H, D)
volume_data[0, 1, 1, 1] = 1
@@ -567,6 +723,11 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume_data, 0.9, return_local_coords=False)
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)
@@ -575,6 +736,12 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
# Check all values are in the range [-1, 1]
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
# test C++ implementation
verts, faces = marching_cubes(volume_data, 0.9, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
def test_cube_no_duplicate_verts(self):
volume_data = torch.zeros(1, 5, 5, 5) # (B, W, H, D)
volume_data[0, 1, 1, 1] = 1
@@ -670,6 +837,11 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
# test C++ implementation
verts, faces = marching_cubes(volume, 64, return_local_coords=False)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
verts, faces = marching_cubes_naive(
volume, isolevel=64, return_local_coords=True
)
@@ -681,6 +853,12 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
# Check all values are in the range [-1, 1]
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
# test C++ implementation
verts, faces = marching_cubes(volume, 64, return_local_coords=True)
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
# Uses skimage.draw.ellipsoid
def test_double_ellipsoid(self):
if USE_SCIKIT:
@@ -694,6 +872,7 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
volume = torch.Tensor(ellip_double).unsqueeze(0)
volume = volume.permute(0, 3, 2, 1) # (B, D, H, W)
verts, faces = marching_cubes_naive(volume, isolevel=0.001)
verts2, faces2 = marching_cubes(volume, isolevel=0.001)
data_filename = "test_marching_cubes_data/double_ellipsoid.pickle"
filename = os.path.join(DATA_DIR, data_filename)
@@ -704,6 +883,8 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
self.assertClose(verts[0], expected_verts)
self.assertClose(faces[0], expected_faces)
self.assertClose(verts2[0], expected_verts)
self.assertClose(faces2[0], expected_faces)
def test_cube_surface_area(self):
if USE_SCIKIT:
@@ -720,12 +901,15 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
volume_data[0, 2, 2, 2] = 1
volume_data = volume_data.permute(0, 3, 2, 1) # (B, D, H, W)
verts, faces = marching_cubes_naive(volume_data, return_local_coords=False)
verts_c, faces_c = marching_cubes(volume_data, return_local_coords=False)
verts_sci, faces_sci = marching_cubes_classic(volume_data[0])
surf = mesh_surface_area(verts[0], faces[0])
surf_c = mesh_surface_area(verts_c[0], faces_c[0])
surf_sci = mesh_surface_area(verts_sci, faces_sci)
self.assertClose(surf, surf_sci)
self.assertClose(surf, surf_c)
def test_sphere_surface_area(self):
if USE_SCIKIT:
@@ -746,12 +930,15 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
).unsqueeze(0)
volume = volume.permute(0, 3, 2, 1) # (B, D, H, W)
verts, faces = marching_cubes_naive(volume, isolevel=64)
verts_c, faces_c = marching_cubes(volume, isolevel=64)
verts_sci, faces_sci = marching_cubes_classic(volume[0], level=64)
surf = mesh_surface_area(verts[0], faces[0])
surf_c = mesh_surface_area(verts_c[0], faces_c[0])
surf_sci = mesh_surface_area(verts_sci, faces_sci)
self.assertClose(surf, surf_sci)
self.assertClose(surf, surf_c)
def test_double_ellipsoid_surface_area(self):
if USE_SCIKIT:
@@ -766,12 +953,15 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
volume = torch.Tensor(ellip_double).unsqueeze(0)
volume = volume.permute(0, 3, 2, 1) # (B, D, H, W)
verts, faces = marching_cubes_naive(volume, isolevel=0)
verts_c, faces_c = marching_cubes(volume, isolevel=0)
verts_sci, faces_sci = marching_cubes_classic(volume[0], level=0)
surf = mesh_surface_area(verts[0], faces[0])
surf_c = mesh_surface_area(verts_c[0], faces_c[0])
surf_sci = mesh_surface_area(verts_sci, faces_sci)
self.assertClose(surf, surf_sci)
self.assertClose(surf, surf_c)
def test_ball_example(self):
N = 15
@@ -780,6 +970,9 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
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)
verts2, faces2 = marching_cubes(u, 0, return_local_coords=False)
self.assertClose(verts[0], verts2[0])
self.assertClose(faces[0], faces2[0])
@staticmethod
def marching_cubes_with_init(algo_type: str, batch_size: int, V: int):
@@ -789,6 +982,7 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
)
algo_table = {
"naive": marching_cubes_naive,
"cextension": marching_cubes,
}
def convert():