pytorch3d/tests/test_marching_cubes.py
Tim Hatch 34bbb3ad32 apply import merging for fbcode/vision/fair (2 of 2)
Summary:
Applies new import merging and sorting from µsort v1.0.

When merging imports, µsort will make a best-effort to move associated
comments to match merged elements, but there are known limitations due to
the diynamic nature of Python and developer tooling. These changes should
not produce any dangerous runtime changes, but may require touch-ups to
satisfy linters and other tooling.

Note that µsort uses case-insensitive, lexicographical sorting, which
results in a different ordering compared to isort. This provides a more
consistent sorting order, matching the case-insensitive order used when
sorting import statements by module name, and ensures that "frog", "FROG",
and "Frog" always sort next to each other.

For details on µsort's sorting and merging semantics, see the user guide:
https://usort.readthedocs.io/en/stable/guide.html#sorting

Reviewed By: bottler

Differential Revision: D35553814

fbshipit-source-id: be49bdb6a4c25264ff8d4db3a601f18736d17be1
2022-04-13 06:51:33 -07:00

775 lines
28 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import os
import pickle
import unittest
import torch
from common_testing import get_tests_dir, TestCaseMixin
from pytorch3d.ops.marching_cubes import marching_cubes_naive
USE_SCIKIT = False
DATA_DIR = get_tests_dir() / "data"
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_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_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