mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
8fe6934885
Summary: Simplify code fixes https://github.com/facebookresearch/pytorch3d/issues/1788 Reviewed By: MichaelRamamonjisoa Differential Revision: D61847675 fbshipit-source-id: 48400875d1d885bb3615bc9f4b3c7c3d822b67e7
244 lines
8.5 KiB
Python
244 lines
8.5 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 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
|