From 47919b4b7e44099cd41585380dc3158576d9c064 Mon Sep 17 00:00:00 2001 From: Stuart Date: Thu, 9 May 2024 11:13:47 -0400 Subject: [PATCH] 3D export fixes and tests (#4368) * WIP 3D export fixes and tests * print * errant copypasta * move rewrite asset logic to scene_3d * cleanup 3d import/export tests a bit --- fiftyone/core/threed/scene_3d.py | 44 ++++ fiftyone/utils/data/exporters.py | 46 +--- tests/unittests/import_export_tests.py | 307 +++++++++++++++++++++++++ 3 files changed, 359 insertions(+), 38 deletions(-) diff --git a/fiftyone/core/threed/scene_3d.py b/fiftyone/core/threed/scene_3d.py index df32d137618..82d94ded6f2 100644 --- a/fiftyone/core/threed/scene_3d.py +++ b/fiftyone/core/threed/scene_3d.py @@ -248,6 +248,10 @@ def traverse(self, include_self=False): Args: include_self: whether to include the current node in the traversal + + Yields: + :class:`Object3D` + """ if include_self: yield self @@ -255,6 +259,46 @@ def traverse(self, include_self=False): for child in self.children: yield from child.traverse(include_self=True) + def update_asset_paths(self, asset_rewrite_paths: dict): + """Update asset paths in this scene according to an input dict mapping. + + Asset path is unchanged if it does not exist in ``asset_rewrite_paths`` + + Args: + asset_rewrite_paths: ``dict`` mapping asset path to new asset path + + Returns: + ``True`` if the scene was modified. + """ + scene_modified = False + for node in self.traverse(): + for path_attribute in node._asset_path_fields: + asset_path = getattr(node, path_attribute, None) + new_asset_path = asset_rewrite_paths.get(asset_path) + + if asset_path is not None and asset_path != new_asset_path: + setattr(node, path_attribute, new_asset_path) + scene_modified = True + + # modify scene background paths, if any + if self.background is not None: + if self.background.image is not None: + new_asset_path = asset_rewrite_paths.get(self.background.image) + if new_asset_path != self.background.image: + self.background.image = new_asset_path + scene_modified = True + + if self.background.cube is not None: + new_cube = [ + asset_rewrite_paths.get(face) + for face in self.background.cube + ] + if new_cube != self.background.cube: + self.background.cube = new_cube + scene_modified = True + + return scene_modified + def get_scene_summary(self): """Returns a summary of the scene.""" node_types = Counter(map(type, self.traverse())) diff --git a/fiftyone/utils/data/exporters.py b/fiftyone/utils/data/exporters.py index 8cb4b85b0cd..3d5ca399fc7 100644 --- a/fiftyone/utils/data/exporters.py +++ b/fiftyone/utils/data/exporters.py @@ -1171,6 +1171,7 @@ def _handle_fo3d_file(self, fo3d_path, fo3d_output_path, export_mode): scene = fo3d.Scene.from_fo3d(fo3d_path) asset_paths = scene.get_asset_paths() + input_to_output_paths = {} for asset_path in asset_paths: if not os.path.isabs(asset_path): absolute_asset_path = os.path.join( @@ -1181,12 +1182,15 @@ def _handle_fo3d_file(self, fo3d_path, fo3d_output_path, export_mode): seen = self._filename_maker.seen_input_path(absolute_asset_path) - if seen: - continue - asset_output_path = self._filename_maker.get_output_path( absolute_asset_path ) + input_to_output_paths[asset_path] = os.path.relpath( + asset_output_path, os.path.dirname(fo3d_output_path) + ) + + if seen: + continue if export_mode is True: etau.copy_file(absolute_asset_path, asset_output_path) @@ -1195,41 +1199,7 @@ def _handle_fo3d_file(self, fo3d_path, fo3d_output_path, export_mode): elif export_mode == "symlink": etau.symlink_file(absolute_asset_path, asset_output_path) - is_scene_modified = False - - for node in scene.traverse(): - path_attribute = next( - ( - attr - for attr in fo3d.fo3d_path_attributes - if hasattr(node, attr) - ), - None, - ) - - if path_attribute is not None: - asset_path = getattr(node, path_attribute) - - is_nested_path = os.path.split(asset_path)[0] != "" - - if asset_path is not None and is_nested_path: - setattr(node, path_attribute, os.path.basename(asset_path)) - is_scene_modified = True - - # modify scene background paths, if any - if scene.background is not None: - if scene.background.image is not None: - scene.background.image = os.path.basename( - scene.background.image - ) - is_scene_modified = True - - if scene.background.cube is not None: - scene.background.cube = [ - os.path.basename(face_path) - for face_path in scene.background.cube - ] - is_scene_modified = True + is_scene_modified = scene.update_asset_paths(input_to_output_paths) if is_scene_modified: # note: we can't have different behavior for "symlink" because diff --git a/tests/unittests/import_export_tests.py b/tests/unittests/import_export_tests.py index 66b14d7df17..eeb0fec9322 100644 --- a/tests/unittests/import_export_tests.py +++ b/tests/unittests/import_export_tests.py @@ -6,8 +6,10 @@ | """ import os +import pathlib import random import string +import tempfile import unittest import cv2 @@ -4648,6 +4650,311 @@ def test_media_directory(self): self.assertEqual(len(relpath.split(os.path.sep)), 2) +class ThreeDMediaTests(unittest.TestCase): + """Tests mostly for proper media export. Labels are tested + properly elsewhere, 3D should be no different in that regard. + """ + + def _build_flat_relative(self, temp_dir): + # Scene has relative asset paths + # Data layout: + # data/ + # image.jpeg + # pcd.pcd + # obj.obj + # mtl.mtl + # s1.fo3d + root_data_dir = os.path.join(temp_dir, "data") + s = fo.Scene() + s.background = fo.SceneBackground(image="image.jpeg") + s.add(fo.PointCloud("pcd", "pcd.pcd")) + s.add(fo.ObjMesh("obj", "obj.obj", "mtl.mtl")) + scene_path = os.path.join(root_data_dir, "s1.fo3d") + s.write(scene_path) + for file in s.get_asset_paths(): + with open(os.path.join(root_data_dir, file), "w") as f: + f.write(file) + dataset = fo.Dataset() + dataset.add_sample(fo.Sample(scene_path)) + return s, dataset + + def _build_flat_absolute(self, temp_dir): + # Scene has absolute asset paths + # Data layout: + # data/ + # image.jpeg + # pcd.pcd + # obj.obj + # mtl.mtl + # s1.fo3d + root_data_dir = os.path.join(temp_dir, "data") + s = fo.Scene() + s.background = fo.SceneBackground( + image=os.path.join(root_data_dir, "image.jpeg") + ) + s.add(fo.PointCloud("pcd", os.path.join(root_data_dir, "pcd.pcd"))) + s.add( + fo.ObjMesh( + "obj", + os.path.join(root_data_dir, "obj.obj"), + os.path.join(root_data_dir, "mtl.mtl"), + ) + ) + scene_path = os.path.join(root_data_dir, "s1.fo3d") + s.write(scene_path) + for file in s.get_asset_paths(): + with open(os.path.join(root_data_dir, file), "w") as f: + f.write(os.path.basename(file)) + + dataset = fo.Dataset() + dataset.add_sample(fo.Sample(scene_path)) + return s, dataset + + def _build_nested_relative(self, temp_dir): + # Scene has relative asset paths + # Data layout: + # data/ + # image.jpeg + # label1/ + # test/ + # s.fo3d + # sub/ + # pcd.pcd + # obj.obj + # mtl.mtl + # label2/ + # test/ + # s.fo3d + # sub/ + # pcd.pcd + # obj.obj + # mtl.mtl + root_data_dir = os.path.join(temp_dir, "data") + scene1_dir = os.path.join(root_data_dir, "label1", "test") + + s = fo.Scene() + s.background = fo.SceneBackground(image="../../image.jpeg") + s.add(fo.PointCloud("pcd", "sub/pcd.pcd")) + s.add( + fo.ObjMesh( + "obj", + "sub/obj.obj", + "sub/mtl.mtl", + ) + ) + scene_path = os.path.join(scene1_dir, "s.fo3d") + s.write(scene_path) + + scene2_dir = os.path.join(root_data_dir, "label2", "test") + + scene_path2 = os.path.join(scene2_dir, "s.fo3d") + + # Scene2 is the same except change something small so we know which + # is which. + s.background.color = "red" + s.write(scene_path2) + + # Write content as filename (with 2 suffix for files from scene 2) + for file in s.get_asset_paths(): + f = pathlib.Path(os.path.join(scene1_dir, file)) + f.parent.mkdir(parents=True, exist_ok=True) + f.write_text(os.path.basename(file)) + + if file.endswith("image.jpeg"): + continue + + f = pathlib.Path(os.path.join(scene2_dir, file)) + f.parent.mkdir(parents=True, exist_ok=True) + f.write_text(os.path.basename(file) + "2") + + dataset = fo.Dataset() + dataset.add_samples([fo.Sample(scene_path), fo.Sample(scene_path2)]) + return dataset + + def _assert_scene_content(self, original_scene, scene, export_dir=None): + self.assertEqual(original_scene, scene) + for file in scene.get_asset_paths(): + if export_dir: + file = os.path.join(export_dir, file) + with open(file) as f: + self.assertEqual(f.read(), os.path.basename(file)) + + @drop_datasets + def test_flat_relative(self): + """Tests a simple flat and relative-addressed scene""" + with tempfile.TemporaryDirectory() as temp_dir: + s, dataset = self._build_flat_relative(temp_dir) + + # Export + export_dir = os.path.join(temp_dir, "export") + dataset.export( + export_dir=export_dir, + dataset_type=fo.types.MediaDirectory, + export_media=True, + ) + + # All files flat in export_dir + fileset = set(os.listdir(export_dir)) + self.assertSetEqual( + fileset, + {"image.jpeg", "pcd.pcd", "obj.obj", "mtl.mtl", "s1.fo3d"}, + ) + + # Same file content + scene2 = fo.Scene.from_fo3d(os.path.join(export_dir, "s1.fo3d")) + self._assert_scene_content(s, scene2, export_dir) + + @drop_datasets + def test_flat_absolute(self): + """Tests a simple flat and absolute-addressed scene""" + with tempfile.TemporaryDirectory() as temp_dir: + s, dataset = self._build_flat_absolute(temp_dir) + + # Export it + export_dir = os.path.join(temp_dir, "export") + dataset.export( + export_dir=export_dir, + dataset_type=fo.types.MediaDirectory, + export_media=True, + ) + + # All files flat in export_dir + fileset = set(os.listdir(export_dir)) + self.assertSetEqual( + fileset, + {"image.jpeg", "pcd.pcd", "obj.obj", "mtl.mtl", "s1.fo3d"}, + ) + + # Write temp scene with resolving relative paths, so we can test + # that scenes are equal if relative paths are resolved + tmp_scene = fo.Scene.from_fo3d(os.path.join(export_dir, "s1.fo3d")) + tmp_scene.write( + os.path.join(export_dir, "test.fo3d"), + resolve_relative_paths=True, + ) + scene2 = fo.Scene.from_fo3d(os.path.join(export_dir, "test.fo3d")) + + self._assert_scene_content(s, scene2) + + @drop_datasets + def test_relative_nested_flatten(self): + """Tests nested structure is flattened to export dir. Will require + rename of duplicate asset file names and change of relative asset path + in fo3d file. + """ + with tempfile.TemporaryDirectory() as temp_dir: + dataset = self._build_nested_relative(temp_dir) + + # Export it and flatten (no rel_dir) + export_dir = os.path.join(temp_dir, "export") + dataset.export( + export_dir=export_dir, + dataset_type=fo.types.MediaDirectory, + export_media=True, + ) + + # Flattening should mean duplicate file names gain a '-2' + fileset = set(os.listdir(export_dir)) + self.assertSetEqual( + fileset, + { + "image.jpeg", + "pcd.pcd", + "obj.obj", + "mtl.mtl", + "s.fo3d", + "image.jpeg", + "pcd-2.pcd", + "obj-2.obj", + "mtl-2.mtl", + "s-2.fo3d", + }, + ) + + # Scene 1 + scene1_2 = fo.Scene.from_fo3d(os.path.join(export_dir, "s.fo3d")) + self.assertSetEqual( + set(scene1_2.get_asset_paths()), + { + "image.jpeg", + "pcd.pcd", + "obj.obj", + "mtl.mtl", + }, + ) + # Scene 2 + scene2_2 = fo.Scene.from_fo3d(os.path.join(export_dir, "s-2.fo3d")) + self.assertSetEqual( + set(scene2_2.get_asset_paths()), + { + "image.jpeg", + "pcd-2.pcd", + "obj-2.obj", + "mtl-2.mtl", + }, + ) + + # Make sure we align on scene number from before - remember, scene2 + # has a red background! Swap if necessary + if scene1_2.background.color == "red": + scene2_2, scene1_2 = scene1_2, scene2_2 + + for file in scene1_2.get_asset_paths(): + with open(os.path.join(export_dir, file)) as f: + self.assertEqual(f.read(), os.path.basename(file)) + + for file in scene2_2.get_asset_paths(): + if file.endswith("image.jpeg"): + continue + with open(os.path.join(export_dir, file)) as f: + self.assertEqual( + f.read(), + os.path.basename(file).replace("-2", "") + "2", + ) + + @drop_datasets + def test_relative_nested_maintain(self): + """Tests nested structure is maintained in export dir. No change in + relative asset paths in fo3d file. + """ + with tempfile.TemporaryDirectory() as temp_dir: + dataset = self._build_nested_relative(temp_dir) + + # Export it - with root data dir as rel_dir + root_data_dir = os.path.join(temp_dir, "data") + export_dir = os.path.join(temp_dir, "export") + + dataset.export( + export_dir=export_dir, + dataset_type=fo.types.MediaDirectory, + export_media=True, + rel_dir=root_data_dir, + ) + + scene1 = fo.Scene.from_fo3d( + os.path.join(export_dir, "label1/test/s.fo3d") + ) + self.assertEqual(scene1.background.image, "../../image.jpeg") + + for file in scene1.get_asset_paths(): + with open(os.path.join(export_dir, "label1/test/", file)) as f: + self.assertEqual(f.read(), os.path.basename(file)) + + scene2 = fo.Scene.from_fo3d( + os.path.join(export_dir, "label2/test/s.fo3d") + ) + self.assertEqual(scene2.background.image, "../../image.jpeg") + + for file in scene2.get_asset_paths(): + with open(os.path.join(export_dir, "label2/test/", file)) as f: + if file.endswith("image.jpeg"): + continue + self.assertEqual( + f.read(), + os.path.basename(file) + "2", + ) + + def _relpath(path, start): # Avoids errors related to symlinks in `/tmp` directories return os.path.relpath(os.path.realpath(path), os.path.realpath(start))