load texture flag

Summary: Add flag for loading textures

Reviewed By: nikhilaravi

Differential Revision: D19664437

fbshipit-source-id: 3cc4e6179df9b7e24efff9e7da3b164253f1d775

pytorch3d/io/obj_io.py

@@ -88,7 +88,7 @@ def _open_file(f):
     return f, new_f
 
 
-def load_obj(f_obj):
+def load_obj(f_obj, load_textures=True):
     """
     Load a mesh and textures from a .obj and .mtl file.
     Currently this handles verts, faces, vertex texture uv coordinates, normals,
@@ -146,6 +146,7 @@ def load_obj(f_obj):
     Args:
         f: A file-like object (with methods read, readline, tell, and seek),
            a pathlib path or a string containing a file name.
+        load_textures: Boolean indicating whether material files are loaded
 
     Returns:
         6-element tuple containing
@@ -201,7 +202,7 @@ def load_obj(f_obj):
         data_dir = os.path.dirname(f_obj)
     f_obj, new_f = _open_file(f_obj)
     try:
-        return _load(f_obj, data_dir)
+        return _load(f_obj, data_dir, load_textures=load_textures)
     finally:
         if new_f:
             f_obj.close()
@@ -273,7 +274,7 @@ def _parse_face(
         faces_materials_idx.append(material_idx)
 
 
-def _load(f_obj, data_dir):
+def _load(f_obj, data_dir, load_textures=True):
     """
     Load a mesh from a file-like object. See load_obj function more details.
     Any material files associated with the obj are expected to be in the
@@ -362,15 +363,16 @@ def _load(f_obj, data_dir):
 
     # Load materials
     material_colors, texture_images = None, None
-    if (len(material_names) > 0) and (f_mtl is not None):
-        if os.path.isfile(f_mtl):
-            material_colors, texture_images = load_mtl(
-                f_mtl, material_names, data_dir
-            )
-        else:
-            warnings.warn(f"Mtl file does not exist: {f_mtl}")
-    elif len(material_names) > 0:
-        warnings.warn("No mtl file provided")
+    if load_textures:
+        if (len(material_names) > 0) and (f_mtl is not None):
+            if os.path.isfile(f_mtl):
+                material_colors, texture_images = load_mtl(
+                    f_mtl, material_names, data_dir
+                )
+            else:
+                warnings.warn(f"Mtl file does not exist: {f_mtl}")
+        elif len(material_names) > 0:
+            warnings.warn("No mtl file provided")
 
     faces = _Faces(
         verts_idx=faces_verts_idx,

tests/test_obj_io.py

@@ -390,6 +390,17 @@ class TestMeshObjIO(unittest.TestCase):
                     )
                 )
 
+    def test_load_mtl_noload(self):
+        DATA_DIR = (
+            Path(__file__).resolve().parent.parent / "docs/tutorials/data"
+        )
+        obj_filename = "cow_mesh/cow.obj"
+        filename = os.path.join(DATA_DIR, obj_filename)
+        verts, faces, aux = load_obj(filename, load_textures=False)
+
+        self.assertTrue(aux.material_colors is None)
+        self.assertTrue(aux.texture_images is None)
+
     def test_load_mtl_fail(self):
         # Faces have a material
         obj_file = "\n".join(
@@ -444,6 +455,27 @@ class TestMeshObjIO(unittest.TestCase):
         self.assertTrue(torch.allclose(verts, expected_verts))
         self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
 
+    def test_load_obj_missing_texture_noload(self):
+        DATA_DIR = Path(__file__).resolve().parent / "data"
+        obj_filename = "missing_files_obj/model.obj"
+        filename = os.path.join(DATA_DIR, obj_filename)
+        verts, faces, aux = load_obj(filename, load_textures=False)
+
+        expected_verts = torch.tensor(
+            [
+                [0.1, 0.2, 0.3],
+                [0.2, 0.3, 0.4],
+                [0.3, 0.4, 0.5],
+                [0.4, 0.5, 0.6],
+            ],
+            dtype=torch.float32,
+        )
+        expected_faces = torch.tensor([[0, 1, 2], [0, 1, 3]], dtype=torch.int64)
+        self.assertTrue(torch.allclose(verts, expected_verts))
+        self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
+        self.assertTrue(aux.material_colors is None)
+        self.assertTrue(aux.texture_images is None)
+
     def test_load_obj_missing_mtl(self):
         DATA_DIR = Path(__file__).resolve().parent / "data"
         obj_filename = "missing_files_obj/model2.obj"
@@ -464,6 +496,27 @@ class TestMeshObjIO(unittest.TestCase):
         self.assertTrue(torch.allclose(verts, expected_verts))
         self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
 
+    def test_load_obj_missing_mtl_noload(self):
+        DATA_DIR = Path(__file__).resolve().parent / "data"
+        obj_filename = "missing_files_obj/model2.obj"
+        filename = os.path.join(DATA_DIR, obj_filename)
+        verts, faces, aux = load_obj(filename, load_textures=False)
+
+        expected_verts = torch.tensor(
+            [
+                [0.1, 0.2, 0.3],
+                [0.2, 0.3, 0.4],
+                [0.3, 0.4, 0.5],
+                [0.4, 0.5, 0.6],
+            ],
+            dtype=torch.float32,
+        )
+        expected_faces = torch.tensor([[0, 1, 2], [0, 1, 3]], dtype=torch.int64)
+        self.assertTrue(torch.allclose(verts, expected_verts))
+        self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
+        self.assertTrue(aux.material_colors is None)
+        self.assertTrue(aux.texture_images is None)
+
     @staticmethod
     def save_obj_with_init(V: int, F: int):
         verts_list = torch.tensor(V * [[0.11, 0.22, 0.33]]).view(-1, 3)