# 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 unittest import torch from pytorch3d.ops.subdivide_meshes import SubdivideMeshes from pytorch3d.structures.meshes import Meshes from pytorch3d.utils.ico_sphere import ico_sphere from .common_testing import TestCaseMixin class TestSubdivideMeshes(TestCaseMixin, unittest.TestCase): def simple_subdivide(self, with_init=False): # Create a mesh with one face and check the subdivided mesh has # 4 faces with the correct vertex coordinates. device = torch.device("cuda:0") verts = torch.tensor( [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=torch.float32, device=device, requires_grad=True, ) faces = torch.tensor([[0, 1, 2]], dtype=torch.int64, device=device) mesh = Meshes(verts=[verts], faces=[faces]) mesh_init = mesh.clone() if with_init else None subdivide = SubdivideMeshes(meshes=mesh_init) new_mesh = subdivide(mesh) # Subdivided face: # # v0 # /\ # / \ # / f0 \ # v4 /______\ v3 # /\ /\ # / \ f3 / \ # / f2 \ / f1 \ # /______\/______\ # v2 v5 v1 # gt_subdivide_verts = torch.tensor( [ [0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.75, 0.5, 0.0], [0.25, 0.5, 0.0], [0.5, 0.0, 0.0], ], dtype=torch.float32, device=device, ) gt_subdivide_faces = torch.tensor( [[0, 3, 4], [1, 5, 3], [2, 4, 5], [5, 4, 3]], dtype=torch.int64, device=device, ) new_verts, new_faces = new_mesh.get_mesh_verts_faces(0) self.assertClose(new_verts, gt_subdivide_verts) self.assertClose(new_faces, gt_subdivide_faces) self.assertTrue(new_verts.requires_grad == verts.requires_grad) def test_simple_subdivide(self): self.simple_subdivide() def test_simple_subdivide_with_init(self): self.simple_subdivide(with_init=True) def test_heterogeneous_meshes(self): device = torch.device("cuda:0") verts1 = torch.tensor( [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=torch.float32, device=device, requires_grad=True, ) faces1 = torch.tensor([[0, 1, 2]], dtype=torch.int64, device=device) verts2 = torch.tensor( [[0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.5, 1.0, 0.0]], dtype=torch.float32, device=device, requires_grad=True, ) faces2 = torch.tensor([[0, 1, 2], [0, 3, 1]], dtype=torch.int64, device=device) faces3 = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int64, device=device) mesh = Meshes(verts=[verts1, verts2, verts2], faces=[faces1, faces2, faces3]) subdivide = SubdivideMeshes() new_mesh = subdivide(mesh.clone()) gt_subdivided_verts1 = torch.tensor( [ [0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.75, 0.5, 0.0], [0.25, 0.5, 0.0], [0.5, 0.0, 0.0], ], dtype=torch.float32, device=device, ) gt_subdivided_faces1 = torch.tensor( [[0, 3, 4], [1, 5, 3], [2, 4, 5], [5, 4, 3]], dtype=torch.int64, device=device, ) # faces2: # # v0 _______e2_______ v3 # /\ / # / \ / # / \ / # e1 / \ e0 / e4 # / \ / # / \ / # / \ / # /______________\/ # v2 e3 v1 # # Subdivided faces2: # # v0 _______v6_______ v3 # /\ /\ / # / \ f1 / \ f3 / # / f0 \ / f7 \ / # v5 /______v4______\/v8 # /\ /\ / # / \ f6 / \ f5 / # / f4 \ / f2 \ / # /______\/______\/ # v2 v7 v1 # gt_subdivided_verts2 = torch.tensor( [ [0.5, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.5, 1.0, 0.0], [0.75, 0.5, 0.0], [0.25, 0.5, 0.0], [1.0, 1.0, 0.0], [0.5, 0.0, 0.0], [1.25, 0.5, 0.0], ], dtype=torch.float32, device=device, ) gt_subdivided_faces2 = torch.tensor( [ [0, 4, 5], [0, 6, 4], [1, 7, 4], [3, 8, 6], [2, 5, 7], [1, 4, 8], [7, 5, 4], [8, 4, 6], ], dtype=torch.int64, device=device, ) gt_subdivided_verts3 = gt_subdivided_verts2.clone() gt_subdivided_verts3[-1, :] = torch.tensor( [0.75, 0.5, 0], dtype=torch.float32, device=device ) gt_subdivided_faces3 = torch.tensor( [ [0, 4, 5], [0, 5, 6], [1, 7, 4], [2, 8, 5], [2, 5, 7], [3, 6, 8], [7, 5, 4], [8, 6, 5], ], dtype=torch.int64, device=device, ) new_mesh_verts1, new_mesh_faces1 = new_mesh.get_mesh_verts_faces(0) new_mesh_verts2, new_mesh_faces2 = new_mesh.get_mesh_verts_faces(1) new_mesh_verts3, new_mesh_faces3 = new_mesh.get_mesh_verts_faces(2) self.assertClose(new_mesh_verts1, gt_subdivided_verts1) self.assertClose(new_mesh_faces1, gt_subdivided_faces1) self.assertClose(new_mesh_verts2, gt_subdivided_verts2) self.assertClose(new_mesh_faces2, gt_subdivided_faces2) self.assertClose(new_mesh_verts3, gt_subdivided_verts3) self.assertClose(new_mesh_faces3, gt_subdivided_faces3) self.assertTrue(new_mesh_verts1.requires_grad == verts1.requires_grad) self.assertTrue(new_mesh_verts2.requires_grad == verts2.requires_grad) self.assertTrue(new_mesh_verts3.requires_grad == verts2.requires_grad) def test_subdivide_features(self): device = torch.device("cuda:0") mesh = ico_sphere(0, device) N = 10 mesh = mesh.extend(N) edges = mesh.edges_packed() V = mesh.num_verts_per_mesh()[0] D = 256 feats = torch.rand( (N * V, D), dtype=torch.float32, device=device, requires_grad=True ) # packed features app_feats = feats[edges].mean(1) subdivide = SubdivideMeshes() new_mesh, new_feats = subdivide(mesh, feats) gt_feats = torch.cat( (feats.view(N, V, D), app_feats.view(N, -1, D)), dim=1 ).view(-1, D) self.assertClose(new_feats, gt_feats) self.assertTrue(new_feats.requires_grad == gt_feats.requires_grad) def test_with_empty(self): verts_list = [[[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], []] faces_list = [[[0, 1, 2], [0, 2, 3]], []] verts_list = [torch.tensor(verts, dtype=torch.float64) for verts in verts_list] face_list = [torch.tensor(faces, dtype=torch.long) for faces in faces_list] meshes = Meshes(verts=verts_list, faces=face_list) subdivided_meshes = SubdivideMeshes()(meshes) self.assertEqual(len(subdivided_meshes), 2) @staticmethod def subdivide_meshes_with_init(num_meshes: int = 10, same_topo: bool = False): device = torch.device("cuda:0") meshes = ico_sphere(0, device=device) if num_meshes > 1: meshes = meshes.extend(num_meshes) meshes_init = meshes.clone() if same_topo else None torch.cuda.synchronize() def subdivide_meshes(): subdivide = SubdivideMeshes(meshes=meshes_init) subdivide(meshes=meshes.clone()) torch.cuda.synchronize() return subdivide_meshes