load texture flag

Summary: Add flag for loading textures

Reviewed By: nikhilaravi

Differential Revision: D19664437

fbshipit-source-id: 3cc4e6179df9b7e24efff9e7da3b164253f1d775
This commit is contained in:
Georgia Gkioxari 2020-01-31 13:35:56 -08:00 committed by Facebook Github Bot
parent 244b7eb80e
commit 659ad34389
2 changed files with 67 additions and 12 deletions

View File

@ -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,

View File

@ -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)