textures dimension check

Summary:
When textures are set on `Meshes` we need to check if the dimensions actually match that of the verts/faces in the mesh. There was a github issue where someone tried to set the attribute after construction of the `Meshes` object and ran into an error when trying to sample textures.

The desired usage is to initialize the class with the textures (not set an attribute afterwards) but in any case we need to check the dimensions match before sampling textures.

Reviewed By: bottler

Differential Revision: D29020249

fbshipit-source-id: 9fb8a5368b83c9ec53652df92b96fc8b2613f591

Author: nikhilaravi

pytorch3d/renderer/mesh/textures.py

@@ -574,6 +574,15 @@ def join_scene(self) -> "TexturesAtlas":
         """
         return self.__class__(atlas=[torch.cat(self.atlas_list())])
 
+    def check_shapes(
+        self, batch_size: int, max_num_verts: int, max_num_faces: int
+    ) -> bool:
+        """
+        Check if the dimensions of the atlas match that of the mesh faces
+        """
+        # (N, F) should be the same
+        return self.atlas_padded().shape[0:2] == (batch_size, max_num_faces)
+
 
 class TexturesUV(TexturesBase):
     def __init__(
@@ -1213,6 +1222,18 @@ def centers_for_image(self, index: int) -> torch.Tensor:
             centers = centers[0, :, 0].T
         return centers
 
+    def check_shapes(
+        self, batch_size: int, max_num_verts: int, max_num_faces: int
+    ) -> bool:
+        """
+        Check if the dimensions of the verts/faces uvs match that of the mesh
+        """
+        # (N, F) should be the same
+        # (N, V) is not guaranteed to be the same
+        return (self.faces_uvs_padded().shape[0:2] == (batch_size, max_num_faces)) and (
+            self.verts_uvs_padded().shape[0] == batch_size
+        )
+
 
 class TexturesVertex(TexturesBase):
     def __init__(
@@ -1292,6 +1313,13 @@ def __getitem__(self, index) -> "TexturesVertex":
         new_props = self._getitem(index, props)
         verts_features = new_props["verts_features_list"]
         if isinstance(verts_features, list):
+            # Handle the case of an empty list
+            if len(verts_features) == 0:
+                verts_features = torch.empty(
+                    size=(0, 0, 3),
+                    dtype=torch.float32,
+                    device=self.verts_features_padded().device,
+                )
             new_tex = self.__class__(verts_features=verts_features)
         elif torch.is_tensor(verts_features):
             new_tex = self.__class__(verts_features=[verts_features])
@@ -1410,3 +1438,12 @@ def join_scene(self) -> "TexturesVertex":
         Return a new TexturesVertex amalgamating the batch.
         """
         return self.__class__(verts_features=[torch.cat(self.verts_features_list())])
+
+    def check_shapes(
+        self, batch_size: int, max_num_verts: int, max_num_faces: int
+    ) -> bool:
+        """
+        Check if the dimensions of the verts features match that of the mesh verts
+        """
+        # (N, V) should be the same
+        return self.verts_features_padded().shape[:-1] == (batch_size, max_num_verts)

pytorch3d/structures/meshes.py

@@ -255,6 +255,7 @@ def __init__(
         if textures is not None and not hasattr(textures, "sample_textures"):
             msg = "Expected textures to be an instance of type TexturesBase; got %r"
             raise ValueError(msg % type(textures))
+
         self.textures = textures
 
         # Indicates whether the meshes in the list/batch have the same number
@@ -424,10 +425,14 @@ def __init__(
             )
 
         # Set the num verts/faces on the textures if present.
-        if self.textures is not None:
+        if textures is not None:
+            shape_ok = self.textures.check_shapes(self._N, self._V, self._F)
+            if not shape_ok:
+                msg = "Textures do not match the dimensions of Meshes."
+                raise ValueError(msg)
+
             self.textures._num_faces_per_mesh = self._num_faces_per_mesh.tolist()
             self.textures._num_verts_per_mesh = self._num_verts_per_mesh.tolist()
-            self.textures._N = self._N
             self.textures.valid = self.valid
 
         if verts_normals is not None:
@@ -1560,6 +1565,13 @@ def extend(self, N: int):
 
     def sample_textures(self, fragments):
         if self.textures is not None:
+
+            # Check dimensions of textures match that of meshes
+            shape_ok = self.textures.check_shapes(self._N, self._V, self._F)
+            if not shape_ok:
+                msg = "Textures do not match the dimensions of Meshes."
+                raise ValueError(msg)
+
             # Pass in faces packed. If the textures are defined per
             # vertex, the face indices are needed in order to interpolate
             # the vertex attributes across the face.

tests/test_texturing.py

@@ -251,7 +251,7 @@ def test_padded_to_packed(self):
     def test_getitem(self):
         N = 5
         V = 20
-        source = {"verts_features": torch.randn(size=(N, 10, 128))}
+        source = {"verts_features": torch.randn(size=(N, V, 128))}
         tex = TexturesVertex(verts_features=source["verts_features"])
 
         verts = torch.rand(size=(N, V, 3))
@@ -268,6 +268,30 @@ def test_getitem(self):
         tryindex(self, index, tex, meshes, source)
         tryindex(self, [2, 4], tex, meshes, source)
 
+    def test_sample_textures_error(self):
+        N = 5
+        V = 20
+        verts = torch.rand(size=(N, V, 3))
+        faces = torch.randint(size=(N, 10, 3), high=V)
+        tex = TexturesVertex(verts_features=torch.randn(size=(N, 10, 128)))
+
+        # Verts features have the wrong number of verts
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Verts features have the wrong batch dim
+        tex = TexturesVertex(verts_features=torch.randn(size=(1, V, 128)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
 
 class TestTexturesAtlas(TestCaseMixin, unittest.TestCase):
     def test_sample_texture_atlas(self):
@@ -456,11 +480,12 @@ def test_padded_to_packed(self):
     def test_getitem(self):
         N = 5
         V = 20
-        source = {"atlas": torch.randn(size=(N, 10, 4, 4, 3))}
+        F = 10
+        source = {"atlas": torch.randn(size=(N, F, 4, 4, 3))}
         tex = TexturesAtlas(atlas=source["atlas"])
 
         verts = torch.rand(size=(N, V, 3))
-        faces = torch.randint(size=(N, 10, 3), high=V)
+        faces = torch.randint(size=(N, F, 3), high=V)
         meshes = Meshes(verts=verts, faces=faces, textures=tex)
 
         tryindex(self, 2, tex, meshes, source)
@@ -473,6 +498,32 @@ def test_getitem(self):
         tryindex(self, index, tex, meshes, source)
         tryindex(self, [2, 4], tex, meshes, source)
 
+    def test_sample_textures_error(self):
+        N = 1
+        V = 20
+        F = 10
+        verts = torch.rand(size=(5, V, 3))
+        faces = torch.randint(size=(5, F, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces)
+
+        # TexturesAtlas have the wrong batch dim
+        tex = TexturesAtlas(atlas=torch.randn(size=(1, F, 4, 4, 3)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # TexturesAtlas have the wrong number of faces
+        tex = TexturesAtlas(atlas=torch.randn(size=(N, 15, 4, 4, 3)))
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
 
 class TestTexturesUV(TestCaseMixin, unittest.TestCase):
     def setUp(self) -> None:
@@ -824,9 +875,10 @@ def test_mesh_to(self):
     def test_getitem(self):
         N = 5
         V = 20
+        F = 10
         source = {
             "maps": torch.rand(size=(N, 1, 1, 3)),
-            "faces_uvs": torch.randint(size=(N, 10, 3), high=V),
+            "faces_uvs": torch.randint(size=(N, F, 3), high=V),
             "verts_uvs": torch.randn(size=(N, V, 2)),
         }
         tex = TexturesUV(
@@ -836,7 +888,7 @@ def test_getitem(self):
         )
 
         verts = torch.rand(size=(N, V, 3))
-        faces = torch.randint(size=(N, 10, 3), high=V)
+        faces = torch.randint(size=(N, F, 3), high=V)
         meshes = Meshes(verts=verts, faces=faces, textures=tex)
 
         tryindex(self, 2, tex, meshes, source)
@@ -858,6 +910,46 @@ def test_centers_for_image(self):
         expected = torch.FloatTensor([[32, 224], [64, 96], [64, 128]])
         self.assertClose(tex.centers_for_image(0), expected)
 
+    def test_sample_textures_error(self):
+        N = 1
+        V = 20
+        F = 10
+        maps = torch.rand(size=(N, 1, 1, 3))
+        verts_uvs = torch.randn(size=(N, V, 2))
+        tex = TexturesUV(
+            maps=maps,
+            faces_uvs=torch.randint(size=(N, 15, 3), high=V),
+            verts_uvs=verts_uvs,
+        )
+        verts = torch.rand(size=(5, V, 3))
+        faces = torch.randint(size=(5, 10, 3), high=V)
+        meshes = Meshes(verts=verts, faces=faces)
+
+        # Wrong number of faces
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Wrong batch dim for faces
+        tex = TexturesUV(
+            maps=maps,
+            faces_uvs=torch.randint(size=(1, F, 3), high=V),
+            verts_uvs=verts_uvs,
+        )
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            Meshes(verts=verts, faces=faces, textures=tex)
+
+        # Wrong batch dim for verts_uvs is not necessary to check as
+        # there is already a check inside TexturesUV for a batch dim
+        # mismatch with faces_uvs
+
+        meshes = Meshes(verts=verts, faces=faces)
+        meshes.textures = tex
+
+        # Cannot use the texture attribute set on meshes for sampling
+        # textures if the dimensions don't match
+        with self.assertRaisesRegex(ValueError, "do not match the dimensions"):
+            meshes.sample_textures(None)
+
 
 class TestRectanglePacking(TestCaseMixin, unittest.TestCase):
     def setUp(self) -> None: