diff --git a/pytorch3d/io/obj_io.py b/pytorch3d/io/obj_io.py index 6a569a30..165097b1 100644 --- a/pytorch3d/io/obj_io.py +++ b/pytorch3d/io/obj_io.py @@ -712,6 +712,7 @@ def save_obj( faces_normals_idx: LongTensor of shape (F, 3) giving the index into normals for each vertex in the face. verts_uvs: FloatTensor of shape (V, 2) giving the uv coordinate per vertex. + When provided with faces_uvs, UVs are saved even without a texture_map. faces_uvs: LongTensor of shape (F, 3) giving the index into verts_uvs for each vertex in the face. texture_map: FloatTensor of shape (H, W, 3) representing the texture map @@ -757,7 +758,8 @@ def save_obj( if path_manager is None: path_manager = PathManager() - save_texture = all(t is not None for t in [faces_uvs, verts_uvs, texture_map]) + save_uvs = faces_uvs is not None and verts_uvs is not None + save_texture = save_uvs and texture_map is not None output_path = Path(f) # Save the .obj file @@ -778,7 +780,7 @@ def save_obj( faces_normals_idx=faces_normals_idx, verts_uvs=verts_uvs, faces_uvs=faces_uvs, - save_texture=save_texture, + save_uvs=save_uvs, save_normals=normals is not None, ) @@ -816,7 +818,7 @@ def _save( faces_normals_idx: Optional[torch.Tensor] = None, verts_uvs: Optional[torch.Tensor] = None, faces_uvs: Optional[torch.Tensor] = None, - save_texture: bool = False, + save_uvs: bool = False, save_normals: bool = False, ) -> None: if len(verts) and (verts.dim() != 2 or verts.size(1) != 3): @@ -851,7 +853,7 @@ def _save( assert faces_normals_idx is not None lines += _write_normals(normals, faces_normals_idx, float_str) - if save_texture: + if save_uvs: assert faces_uvs is not None assert verts_uvs is not None @@ -881,7 +883,7 @@ def _save( _write_faces( f, faces, - faces_uvs if save_texture else None, + faces_uvs if save_uvs else None, faces_normals_idx if save_normals else None, ) diff --git a/tests/test_io_obj.py b/tests/test_io_obj.py index 3668be4e..eb379ea5 100644 --- a/tests/test_io_obj.py +++ b/tests/test_io_obj.py @@ -955,6 +955,60 @@ class TestMeshObjIO(TestCaseMixin, unittest.TestCase): with open(obj_file, "r") as actual_file: self.assertEqual(actual_file.read(), expected_obj_file) + def test_save_obj_with_uvs_without_texture_map(self): + verts = torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]], + dtype=torch.float32, + ) + faces = torch.tensor([[0, 2, 1]], dtype=torch.int64) + verts_uvs = torch.tensor( + [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]], + dtype=torch.float32, + ) + faces_uvs = torch.tensor([[3, 1, 2]], dtype=torch.int64) + normals = torch.tensor([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]]) + faces_normals_idx = torch.tensor([[1, 0, 1]], dtype=torch.int64) + + for save_normals in (False, True): + with ( + self.subTest(save_normals=save_normals), + TemporaryDirectory() as temp_dir, + ): + obj_file = Path(temp_dir) / "mesh.obj" + save_obj( + obj_file, + verts, + faces, + verts_uvs=verts_uvs, + faces_uvs=faces_uvs, + normals=normals if save_normals else None, + faces_normals_idx=faces_normals_idx if save_normals else None, + ) + + self.assertEqual(os.listdir(temp_dir), ["mesh.obj"]) + obj_text = obj_file.read_text() + self.assertNotIn("mtllib", obj_text) + self.assertNotIn("usemtl", obj_text) + expected_face = ( + "f 1/4/2 3/2/1 2/3/2" if save_normals else "f 1/4 3/2 2/3" + ) + self.assertIn(expected_face, obj_text) + + loaded_verts, loaded_faces, aux = load_obj( + obj_file, load_textures=False + ) + self.assertClose(loaded_verts, verts) + self.assertClose(loaded_faces.verts_idx, faces) + self.assertClose(aux.verts_uvs, verts_uvs) + self.assertClose(loaded_faces.textures_idx, faces_uvs) + self.assertIsNone(aux.material_colors) + self.assertIsNone(aux.texture_images) + if save_normals: + self.assertClose(aux.normals, normals) + self.assertClose(loaded_faces.normals_idx, faces_normals_idx) + else: + self.assertIsNone(aux.normals) + def test_save_obj_with_texture(self): verts = torch.tensor( [[0.01, 0.2, 0.301], [0.2, 0.03, 0.408], [0.3, 0.4, 0.05], [0.6, 0.7, 0.8]],