Skip to content

Commit e46ab49

Browse files
ttanpcsfacebook-github-bot
authored andcommitted
Submeshing TexturesAtlas for PyTorch3D 3D Rendering
Summary: Implement submeshing for TexturesAtlas and add associated test Reviewed By: bottler Differential Revision: D52334053 fbshipit-source-id: d54080e9af1f0c01551702736e858e3bd439ac58
1 parent 8a27590 commit e46ab49

File tree

3 files changed

+60
-2
lines changed

3 files changed

+60
-2
lines changed

pytorch3d/renderer/mesh/textures.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -549,6 +549,33 @@ def sample_textures(self, fragments, **kwargs) -> torch.Tensor:
549549

550550
return texels
551551

552+
def submeshes(
553+
self,
554+
vertex_ids_list: List[List[torch.LongTensor]],
555+
faces_ids_list: List[List[torch.LongTensor]],
556+
) -> "TexturesAtlas":
557+
"""
558+
Extract a sub-texture for use in a submesh.
559+
560+
If the meshes batch corresponding to this TextureAtlas contains
561+
`n = len(faces_ids_list)` meshes, then self.atlas_list()
562+
will be of length n. After submeshing, we obtain a batch of
563+
`k = sum(len(v) for v in atlas_list` submeshes (see Meshes.submeshes). This
564+
function creates a corresponding TexturesAtlas object with `atlas_list`
565+
of length `k`.
566+
"""
567+
if len(faces_ids_list) != len(self.atlas_list()):
568+
raise IndexError(
569+
"faces_ids_list must be of " "the same length as atlas_list."
570+
)
571+
572+
sub_features = []
573+
for atlas, faces_ids in zip(self.atlas_list(), faces_ids_list):
574+
for faces_ids_submesh in faces_ids:
575+
sub_features.append(atlas[faces_ids_submesh])
576+
577+
return self.__class__(sub_features)
578+
552579
def faces_verts_textures_packed(self) -> torch.Tensor:
553580
"""
554581
Samples texture from each vertex for each face in the mesh.

pytorch3d/structures/meshes.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1576,8 +1576,6 @@ def submeshes(
15761576
Returns:
15771577
Meshes object of length `sum(len(ids) for ids in face_indices)`.
15781578
1579-
Submeshing only works with no textures, TexturesVertex, or TexturesUV.
1580-
15811579
Example 1:
15821580
15831581
If `meshes` has batch size 1, and `face_indices` is a 1D LongTensor,

tests/test_texturing.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -576,6 +576,39 @@ def test_sample_textures_error(self):
576576
with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
577577
meshes.sample_textures(None)
578578

579+
def test_submeshes(self):
580+
N = 2
581+
V = 5
582+
F = 5
583+
tex = TexturesAtlas(
584+
atlas=torch.arange(N * F * 4 * 4 * 3, dtype=torch.float32).reshape(
585+
N, F, 4, 4, 3
586+
)
587+
)
588+
589+
verts = torch.rand(size=(N, V, 3))
590+
faces = torch.randint(size=(N, F, 3), high=V)
591+
mesh = Meshes(verts=verts, faces=faces, textures=tex)
592+
593+
sub_faces = [
594+
[torch.tensor([0, 2]), torch.tensor([1, 2])],
595+
[],
596+
]
597+
subtex = mesh.submeshes(sub_faces).textures
598+
subtex_faces = subtex.atlas_list()
599+
600+
self.assertEqual(len(subtex_faces), 2)
601+
self.assertClose(
602+
subtex_faces[0].flatten().msort(),
603+
torch.cat(
604+
(
605+
torch.arange(4 * 4 * 3, dtype=torch.float32),
606+
torch.arange(96, 96 + 4 * 4 * 3, dtype=torch.float32),
607+
),
608+
0,
609+
),
610+
)
611+
579612

580613
class TestTexturesUV(TestCaseMixin, unittest.TestCase):
581614
def setUp(self) -> None:

0 commit comments

Comments
 (0)