mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-12-19 22:00:35 +08:00
Classic Marching Cubes algorithm implementation
Summary: Defines a function to run marching cubes algorithm on a single or batch of 3D scalar fields. Returns a mesh's faces and vertices. UPDATES (12/18) - Input data is now specified as a (B, D, H, W) tensor as opposed to a (B, W, H, D) tensor. This will now be compatible with the Volumes datastructure. - Add an option to return output vertices in local coordinates instead of world coordinates. Also added a small fix to remove the dype for device in Transforms3D - if passing in a torch.device instead of str it causes a pyre error. Reviewed By: jcjohnson Differential Revision: D24599019 fbshipit-source-id: 90554a200319fed8736a12371cc349e7108aacd0
This commit is contained in:
Nikhila Ravi
committed by
Facebook GitHub Bot
Facebook GitHub Bot
parent
9c6b58c5ad
commit
ebac66daeb
25
tests/bm_marching_cubes.py
Normal file
25
tests/bm_marching_cubes.py
Normal file
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
from fvcore.common.benchmark import benchmark
|
||||
from test_marching_cubes import TestMarchingCubes
|
||||
|
||||
|
||||
def bm_marching_cubes() -> None:
|
||||
kwargs_list = [
|
||||
{"batch_size": 1, "V": 5},
|
||||
{"batch_size": 1, "V": 10},
|
||||
{"batch_size": 1, "V": 20},
|
||||
{"batch_size": 1, "V": 40},
|
||||
{"batch_size": 5, "V": 5},
|
||||
{"batch_size": 20, "V": 20},
|
||||
]
|
||||
benchmark(
|
||||
TestMarchingCubes.marching_cubes_with_init,
|
||||
"MARCHING_CUBES",
|
||||
kwargs_list,
|
||||
warmup_iters=1,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
bm_marching_cubes()
|
||||
BIN
tests/data/test_marching_cubes_data/double_ellipsoid.pickle
Normal file
BIN
tests/data/test_marching_cubes_data/double_ellipsoid.pickle
Normal file
Binary file not shown.
BIN
tests/data/test_marching_cubes_data/sphere_level64.pickle
Normal file
BIN
tests/data/test_marching_cubes_data/sphere_level64.pickle
Normal file
Binary file not shown.
771
tests/test_marching_cubes.py
Normal file
771
tests/test_marching_cubes.py
Normal file
@@ -0,0 +1,771 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
import os
|
||||
import pickle
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from common_testing import TestCaseMixin
|
||||
from pytorch3d.ops.marching_cubes import marching_cubes_naive
|
||||
|
||||
|
||||
USE_SCIKIT = False
|
||||
|
||||
|
||||
def convert_to_local(verts, volume_dim):
|
||||
return (2 * verts) / (volume_dim - 1) - 1
|
||||
|
||||
|
||||
class TestCubeConfiguration(TestCaseMixin, unittest.TestCase):
|
||||
|
||||
# Test single cubes. Each case corresponds to the corresponding
|
||||
# cube vertex configuration in each case here (0-indexed):
|
||||
# https://en.wikipedia.org/wiki/Marching_cubes#/media/File:MarchingCubes.svg
|
||||
|
||||
def test_empty_volume(self): # case 0
|
||||
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)
|
||||
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
|
||||
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],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[1, 2, 0]])
|
||||
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)
|
||||
|
||||
def test_case2(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0:2, 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(
|
||||
[
|
||||
[1.0000, 0.0000, 0.5000],
|
||||
[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]])
|
||||
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)
|
||||
|
||||
def test_case3(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 1, 1, 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.5000, 0.0000, 0.0000],
|
||||
[0.0000, 0.0000, 0.5000],
|
||||
[1.0000, 1.0000, 0.5000],
|
||||
[0.5000, 1.0000, 0.0000],
|
||||
[1.0000, 0.5000, 0.0000],
|
||||
[0.0000, 0.5000, 0.0000],
|
||||
]
|
||||
)
|
||||
expected_faces = torch.tensor([[0, 1, 5], [4, 3, 2]])
|
||||
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)
|
||||
|
||||
def test_case4(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 1, 0, 0] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 0, 0, 1] = 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.5000, 0.0000, 0.0000],
|
||||
[0.0000, 0.0000, 0.5000],
|
||||
[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]])
|
||||
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)
|
||||
|
||||
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
|
||||
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.0000, 0.5000, 1.0000],
|
||||
[1.0000, 0.5000, 1.0000],
|
||||
[1.0000, 0.5000, 0.0000],
|
||||
[0.0000, 0.5000, 0.0000],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[1, 0, 2], [2, 0, 3]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case6(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 1, 0, 0] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 0, 0, 1] = 0
|
||||
volume_data[0, 0, 1, 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.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, 1.0000],
|
||||
[1.0000, 0.5000, 1.0000],
|
||||
[1.0000, 0.5000, 0.0000],
|
||||
[0.0000, 0.5000, 0.0000],
|
||||
]
|
||||
)
|
||||
expected_faces = torch.tensor([[2, 7, 3], [0, 6, 1], [6, 4, 1], [6, 5, 4]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case7(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 1, 1, 0] = 0
|
||||
volume_data[0, 0, 1, 1] = 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.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],
|
||||
[1.0000, 0.5000, 1.0000],
|
||||
[1.0000, 0.5000, 0.0000],
|
||||
[0.0000, 0.5000, 0.0000],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[0, 1, 9], [4, 7, 8], [2, 3, 11], [5, 10, 6]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case8(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 0, 0, 1] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 0, 1, 1] = 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(
|
||||
[
|
||||
[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, 0.5000, 0.0000],
|
||||
]
|
||||
)
|
||||
expected_faces = torch.tensor([[2, 3, 5], [4, 2, 5], [4, 5, 1], [4, 1, 0]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case9(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 1, 0, 0] = 0
|
||||
volume_data[0, 0, 0, 1] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 0, 1, 1] = 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.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],
|
||||
]
|
||||
)
|
||||
expected_faces = torch.tensor([[0, 5, 4], [0, 4, 3], [0, 3, 1], [3, 4, 2]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case10(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 1, 1, 1] = 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.5000, 0.0000, 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],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[4, 3, 2], [0, 1, 5]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case11(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 1, 0, 0] = 0
|
||||
volume_data[0, 1, 1, 1] = 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(
|
||||
[
|
||||
[1.0000, 0.0000, 0.5000],
|
||||
[0.0000, 0.0000, 0.5000],
|
||||
[0.5000, 1.0000, 1.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],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[5, 1, 6], [5, 0, 1], [4, 3, 2]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case12(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 1, 0, 0] = 0
|
||||
volume_data[0, 0, 1, 0] = 0
|
||||
volume_data[0, 1, 1, 1] = 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(
|
||||
[
|
||||
[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.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]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case13(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 0, 1, 0] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 1, 1, 1] = 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.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],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[3, 6, 2], [3, 7, 6], [1, 5, 0], [5, 4, 0]])
|
||||
|
||||
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)
|
||||
|
||||
def test_case14(self):
|
||||
volume_data = torch.ones(1, 2, 2, 2) # (B, W, H, D)
|
||||
volume_data[0, 0, 0, 0] = 0
|
||||
volume_data[0, 0, 0, 1] = 0
|
||||
volume_data[0, 1, 0, 1] = 0
|
||||
volume_data[0, 1, 1, 1] = 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(
|
||||
[
|
||||
[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],
|
||||
]
|
||||
)
|
||||
|
||||
expected_faces = torch.tensor([[1, 0, 3], [1, 3, 4], [1, 4, 5], [2, 4, 3]])
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
|
||||
def test_single_point(self):
|
||||
volume_data = torch.zeros(1, 3, 3, 3) # (B, W, H, D)
|
||||
volume_data[0, 1, 1, 1] = 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)
|
||||
|
||||
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],
|
||||
]
|
||||
)
|
||||
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],
|
||||
[3, 5, 4],
|
||||
]
|
||||
)
|
||||
|
||||
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())
|
||||
|
||||
def test_cube(self):
|
||||
volume_data = torch.zeros(1, 5, 5, 5) # (B, W, H, D)
|
||||
volume_data[0, 1, 1, 1] = 1
|
||||
volume_data[0, 1, 1, 2] = 1
|
||||
volume_data[0, 2, 1, 1] = 1
|
||||
volume_data[0, 2, 1, 2] = 1
|
||||
volume_data[0, 1, 2, 1] = 1
|
||||
volume_data[0, 1, 2, 2] = 1
|
||||
volume_data[0, 2, 2, 1] = 1
|
||||
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, 0.9, return_local_coords=False)
|
||||
|
||||
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],
|
||||
[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.1000, 1.0000, 1.0000],
|
||||
[2.1000, 1.0000, 2.0000],
|
||||
[2.1000, 2.0000, 1.0000],
|
||||
[2.1000, 2.0000, 2.0000],
|
||||
]
|
||||
)
|
||||
|
||||
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],
|
||||
[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],
|
||||
[9, 17, 19],
|
||||
[11, 9, 19],
|
||||
[20, 13, 12],
|
||||
[20, 21, 14],
|
||||
[13, 20, 14],
|
||||
[15, 14, 21],
|
||||
[22, 20, 12],
|
||||
[16, 22, 12],
|
||||
[21, 20, 23],
|
||||
[23, 20, 22],
|
||||
[17, 15, 21],
|
||||
[23, 17, 21],
|
||||
[22, 16, 18],
|
||||
[23, 22, 18],
|
||||
[19, 23, 18],
|
||||
[17, 23, 19],
|
||||
]
|
||||
)
|
||||
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)
|
||||
self.assertClose(faces[0], expected_faces)
|
||||
|
||||
# Check all values are in the range [-1, 1]
|
||||
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
|
||||
volume_data[0, 1, 1, 2] = 1
|
||||
volume_data[0, 2, 1, 1] = 1
|
||||
volume_data[0, 2, 1, 2] = 1
|
||||
volume_data[0, 1, 2, 1] = 1
|
||||
volume_data[0, 1, 2, 2] = 1
|
||||
volume_data[0, 2, 2, 1] = 1
|
||||
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, 1, return_local_coords=False)
|
||||
|
||||
expected_verts = torch.tensor(
|
||||
[
|
||||
[1.0, 1.0, 1.0],
|
||||
[1.0, 1.0, 2.0],
|
||||
[1.0, 2.0, 1.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],
|
||||
]
|
||||
)
|
||||
|
||||
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],
|
||||
]
|
||||
)
|
||||
|
||||
self.assertClose(verts[0], expected_verts)
|
||||
self.assertClose(faces[0], expected_faces)
|
||||
|
||||
verts, faces = marching_cubes_naive(volume_data, 1, return_local_coords=True)
|
||||
expected_verts = convert_to_local(expected_verts, 5)
|
||||
self.assertClose(verts[0], expected_verts)
|
||||
self.assertClose(faces[0], expected_faces)
|
||||
|
||||
# Check all values are in the range [-1, 1]
|
||||
self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
|
||||
|
||||
def test_sphere(self):
|
||||
# (B, W, H, D)
|
||||
volume = torch.Tensor(
|
||||
[
|
||||
[
|
||||
[(x - 10) ** 2 + (y - 10) ** 2 + (z - 10) ** 2 for z in range(20)]
|
||||
for y in range(20)
|
||||
]
|
||||
for x in range(20)
|
||||
]
|
||||
).unsqueeze(0)
|
||||
volume = volume.permute(0, 3, 2, 1) # (B, D, H, W)
|
||||
verts, faces = marching_cubes_naive(
|
||||
volume, isolevel=64, return_local_coords=False
|
||||
)
|
||||
|
||||
DATA_DIR = Path(__file__).resolve().parent / "data"
|
||||
data_filename = "test_marching_cubes_data/sphere_level64.pickle"
|
||||
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()
|
||||
|
||||
self.assertClose(verts[0], expected_verts)
|
||||
self.assertClose(faces[0], expected_faces)
|
||||
|
||||
verts, faces = marching_cubes_naive(
|
||||
volume, isolevel=64, return_local_coords=True
|
||||
)
|
||||
|
||||
expected_verts = convert_to_local(expected_verts, 20)
|
||||
self.assertClose(verts[0], expected_verts)
|
||||
self.assertClose(faces[0], expected_faces)
|
||||
|
||||
# Check all values are in the range [-1, 1]
|
||||
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:
|
||||
import numpy as np
|
||||
from skimage.draw import ellipsoid
|
||||
|
||||
ellip_base = ellipsoid(6, 10, 16, levelset=True)
|
||||
ellip_double = np.concatenate(
|
||||
(ellip_base[:-1, ...], ellip_base[2:, ...]), axis=0
|
||||
)
|
||||
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)
|
||||
|
||||
DATA_DIR = Path(__file__).resolve().parent / "data"
|
||||
data_filename = "test_marching_cubes_data/double_ellipsoid.pickle"
|
||||
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"]
|
||||
expected_faces = verts_and_faces["faces"]
|
||||
|
||||
self.assertClose(verts[0], expected_verts[0])
|
||||
self.assertClose(faces[0], expected_faces[0])
|
||||
|
||||
def test_cube_surface_area(self):
|
||||
if USE_SCIKIT:
|
||||
from skimage.measure import marching_cubes_classic, mesh_surface_area
|
||||
|
||||
volume_data = torch.zeros(1, 5, 5, 5)
|
||||
volume_data[0, 1, 1, 1] = 1
|
||||
volume_data[0, 1, 1, 2] = 1
|
||||
volume_data[0, 2, 1, 1] = 1
|
||||
volume_data[0, 2, 1, 2] = 1
|
||||
volume_data[0, 1, 2, 1] = 1
|
||||
volume_data[0, 1, 2, 2] = 1
|
||||
volume_data[0, 2, 2, 1] = 1
|
||||
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_sci, faces_sci = marching_cubes_classic(volume_data[0])
|
||||
|
||||
surf = mesh_surface_area(verts[0], faces[0])
|
||||
surf_sci = mesh_surface_area(verts_sci, faces_sci)
|
||||
|
||||
self.assertClose(surf, surf_sci)
|
||||
|
||||
def test_sphere_surface_area(self):
|
||||
if USE_SCIKIT:
|
||||
from skimage.measure import marching_cubes_classic, mesh_surface_area
|
||||
|
||||
# (B, W, H, D)
|
||||
volume = torch.Tensor(
|
||||
[
|
||||
[
|
||||
[
|
||||
(x - 10) ** 2 + (y - 10) ** 2 + (z - 10) ** 2
|
||||
for z in range(20)
|
||||
]
|
||||
for y in range(20)
|
||||
]
|
||||
for x in range(20)
|
||||
]
|
||||
).unsqueeze(0)
|
||||
volume = volume.permute(0, 3, 2, 1) # (B, D, H, W)
|
||||
verts, faces = marching_cubes_naive(volume, isolevel=64)
|
||||
verts_sci, faces_sci = marching_cubes_classic(volume[0], level=64)
|
||||
|
||||
surf = mesh_surface_area(verts[0], faces[0])
|
||||
surf_sci = mesh_surface_area(verts_sci, faces_sci)
|
||||
|
||||
self.assertClose(surf, surf_sci)
|
||||
|
||||
def test_double_ellipsoid_surface_area(self):
|
||||
if USE_SCIKIT:
|
||||
import numpy as np
|
||||
from skimage.draw import ellipsoid
|
||||
from skimage.measure import marching_cubes_classic, mesh_surface_area
|
||||
|
||||
ellip_base = ellipsoid(6, 10, 16, levelset=True)
|
||||
ellip_double = np.concatenate(
|
||||
(ellip_base[:-1, ...], ellip_base[2:, ...]), axis=0
|
||||
)
|
||||
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_sci, faces_sci = marching_cubes_classic(volume[0], level=0)
|
||||
|
||||
surf = mesh_surface_area(verts[0], faces[0])
|
||||
surf_sci = mesh_surface_area(verts_sci, faces_sci)
|
||||
|
||||
self.assertClose(surf, surf_sci)
|
||||
|
||||
@staticmethod
|
||||
def marching_cubes_with_init(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()
|
||||
|
||||
def convert():
|
||||
marching_cubes_naive(volume_data, return_local_coords=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return convert
|
||||
Reference in New Issue
Block a user