From 659ad34389d47bf13f37a340a18a5784ad9b2695 Mon Sep 17 00:00:00 2001 From: Georgia Gkioxari Date: Fri, 31 Jan 2020 13:35:56 -0800 Subject: [PATCH] load texture flag Summary: Add flag for loading textures Reviewed By: nikhilaravi Differential Revision: D19664437 fbshipit-source-id: 3cc4e6179df9b7e24efff9e7da3b164253f1d775 --- pytorch3d/io/obj_io.py | 26 +++++++++++---------- tests/test_obj_io.py | 53 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 12 deletions(-) diff --git a/pytorch3d/io/obj_io.py b/pytorch3d/io/obj_io.py index 1f9aa4aa..d85fa14e 100644 --- a/pytorch3d/io/obj_io.py +++ b/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, diff --git a/tests/test_obj_io.py b/tests/test_obj_io.py index a02ebd80..336dbfdb 100644 --- a/tests/test_obj_io.py +++ b/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)