Make extend respect subclasses of textures

Gil Moshayof · facebook-github-bot · commit 34163326b268 · 2021-05-07T03:42:45.000-07:00
Summary: 3 extend functions in textures.py updated to call `self.__class__` rather than create a new object of its type. As mentioned in #618 . Reviewed By: bottler Differential Revision: D28281218 fbshipit-source-id: b9c99ab87e46a3f28c37efa1ee2c2dceb560b491
diff --git a/pytorch3d/renderer/mesh/textures.py b/pytorch3d/renderer/mesh/textures.py
@@ -470,7 +470,7 @@ def atlas_packed(self) -> torch.Tensor:
 
     def extend(self, N: int) -> "TexturesAtlas":
         new_props = self._extend(N, ["atlas_padded", "_num_faces_per_mesh"])
-        new_tex = TexturesAtlas(atlas=new_props["atlas_padded"])
+        new_tex = self.__class__(atlas=new_props["atlas_padded"])
         new_tex._num_faces_per_mesh = new_props["_num_faces_per_mesh"]
         return new_tex
 
@@ -865,7 +865,7 @@ def extend(self, N: int) -> "TexturesUV":
                 "_num_faces_per_mesh",
             ],
         )
-        new_tex = TexturesUV(
+        new_tex = self.__class__(
             maps=new_props["maps_padded"],
             faces_uvs=new_props["faces_uvs_padded"],
             verts_uvs=new_props["verts_uvs_padded"],
@@ -1339,7 +1339,7 @@ def verts_features_packed(self) -> torch.Tensor:
 
     def extend(self, N: int) -> "TexturesVertex":
         new_props = self._extend(N, ["verts_features_padded", "_num_verts_per_mesh"])
-        new_tex = TexturesVertex(verts_features=new_props["verts_features_padded"])
+        new_tex = self.__class__(verts_features=new_props["verts_features_padded"])
         new_tex._num_verts_per_mesh = new_props["_num_verts_per_mesh"]
         return new_tex
 
