3D compute metadata assets (#4314)

* 3dutils - clean + get_scene_asset_paths

* update compute_metadata for 3d: SceneMetadata

* refactor get_scene_asset_paths to use fos.run to work in batches

* get_scene_asset_paths test

* SceneMetadata.build_for test

* resolve relative asset paths before computing metadata, allowing e.g. "../blah.jpg"

* coderabbit: remove unused import

* coderabbit: improve warning log

* fos.resolve after join for abs paths

* resolve->join_resolve

* Revert "resolve->join_resolve"

This reverts commit 9b34206.

* always resolve if we want abs paths

* minor PR comments

fiftyone/core/metadata.py

@@ -5,22 +5,25 @@
 | `voxel51.com <https://voxel51.com/>`_
 |
 """
-
-import itertools
+import collections
+import json
 import logging
 import multiprocessing.dummy
 import os
+import pathlib
+
 import requests
 
 from PIL import Image
 
 import eta.core.utils as etau
 import eta.core.video as etav
 
-import fiftyone as fo
 from fiftyone.core.odm import DynamicEmbeddedDocument
 import fiftyone.core.fields as fof
 import fiftyone.core.media as fom
+import fiftyone.core.storage as fos
+import fiftyone.core.threed as fo3d
 import fiftyone.core.utils as fou
 
 
@@ -75,6 +78,71 @@ def _build_for_url(cls, url, mime_type=None):
         return cls(size_bytes=size_bytes, mime_type=mime_type)
 
 
+class SceneMetadata(Metadata):
+    """Class for storing metadata about 3D scene samples.
+
+    Args:
+        size_bytes (None): the size of scene definition and all children
+            assets on disk, in bytes
+        mime_type (None): the MIME type of the scene media. Always set to
+            application/octet-stream
+        asset_counts (None): dict of child asset file type to count
+    """
+
+    asset_counts = fof.DictField
+
+    @classmethod
+    def build_for(cls, scene_path, mime_type=None):
+        """Builds a :class:`SceneMetadata` object for the given 3D Scene.
+
+        Args:
+            scene_path: a scene path or URL
+            mime_type (None): Ignored. mime_type always set to
+                application/octet-stream
+
+        Returns:
+            a :class:`SceneMetadata`
+        """
+        if not etau.is_str(scene_path):
+            raise ValueError("Invalid scene or path:", scene_path)
+
+        scene = fo3d.Scene.from_fo3d(scene_path)
+        scene_dict = scene.as_dict()
+        total_size = len(json.dumps(scene_dict))
+        asset_counts = collections.defaultdict(int)
+        mime_type = "application/octet-stream"
+
+        # Get asset paths and resolve all to absolute
+        asset_paths = scene.get_asset_paths()
+        scene_dir = os.path.dirname(scene_path)
+        for i, asset_path in enumerate(asset_paths):
+            if not fos.isabs(asset_path):
+                asset_path = fos.join(scene_dir, asset_path)
+            asset_paths[i] = fos.resolve(asset_path)
+            file_type = pathlib.Path(asset_path).suffix[1:]
+            asset_counts[file_type] += 1
+
+        # Dedupe asset paths within this single scene
+        asset_paths = list(set(asset_paths))
+
+        # compute metadata for all asset paths
+        asset_metadatas = fos.run(
+            _do_compute_metadata,
+            [
+                (i, asset_path, fom.MIXED)
+                for i, asset_path in enumerate(asset_paths)
+            ],
+            progress=False,
+        )
+        total_size += sum(m[1].size_bytes for m in asset_metadatas)
+
+        return cls(
+            size_bytes=total_size,
+            mime_type=mime_type,
+            asset_counts=asset_counts,
+        )
+
+
 class ImageMetadata(Metadata):
     """Class for storing metadata about image samples.
 
@@ -456,6 +524,8 @@ def _get_metadata(filepath, media_type):
         metadata = ImageMetadata.build_for(filepath)
     elif media_type == fom.VIDEO:
         metadata = VideoMetadata.build_for(filepath)
+    elif media_type == fom.THREE_D:
+        metadata = SceneMetadata.build_for(filepath)
     else:
         metadata = Metadata.build_for(filepath)
 

fiftyone/core/storage.py

@@ -281,6 +281,19 @@ def join(a, *p):
     return os.path.join(a, *p)
 
 
+def resolve(path):
+    """Resolves path to absolute, resolving symlinks and relative path
+        indicators such as `.` and `..`.
+
+    Args:
+        path: the filepath
+
+    Returns:
+        the resolved path
+    """
+    return os.path.realpath(path)
+
+
 def isabs(path):
     """Determines whether the given path is absolute.
 

fiftyone/utils/utils3d.py

@@ -6,7 +6,7 @@
 |
 """
 import contextlib
-import itertools
+import functools
 import logging
 import os
 import warnings
@@ -439,35 +439,84 @@ class OrthographicProjectionMetadata(DynamicEmbeddedDocument, fol._HasMedia):
     height = fof.IntField()
 
 
-def _get_pcd_filepath_from_fo3d_scene(scene: Scene, scene_path: str):
-    explicitly_flagged_pcd_path = None
-    fallover_pcd_path = None
+def _get_scene_paths(scene_paths):
+    """Return Tuple of scene paths to use and whether all are local
+    This function is a no-op here but could be different in a repo fork.
+    """
+    return scene_paths, True
 
-    def _visit_node_dfs(node):
-        nonlocal explicitly_flagged_pcd_path
-        nonlocal fallover_pcd_path
 
-        if hasattr(node, "pcd_path") and node.flag_for_projection:
-            explicitly_flagged_pcd_path = node.pcd_path
-        else:
-            if hasattr(node, "pcd_path"):
-                fallover_pcd_path = node.pcd_path
+def _get_scene_asset_paths_single(task, abs_paths=False, skip_failures=True):
+    scene_path, original_scene_path = task
 
-            for child in node.children:
-                _visit_node_dfs(child)
+    # Read scene file which is JSON
+    try:
+        scene = Scene.from_fo3d(scene_path)
+    except Exception as e:
+        if not skip_failures:
+            raise
 
-    _visit_node_dfs(scene)
+        if skip_failures != "ignore":
+            logger.warning(
+                "Failed to process scene at '%s': %s", original_scene_path, e
+            )
+        return []
 
-    pcd_path = (
-        explicitly_flagged_pcd_path
-        if explicitly_flagged_pcd_path
-        else fallover_pcd_path
+    asset_paths = scene.get_asset_paths()
+
+    if abs_paths:
+        # Convert any relative-to-scene paths to absolute
+        scene_dir = os.path.dirname(original_scene_path)
+        for i, asset_path in enumerate(asset_paths):
+            if not fos.isabs(asset_path):
+                asset_path = fos.join(scene_dir, asset_path)
+            asset_paths[i] = fos.resolve(asset_path)
+
+    return asset_paths
+
+
+def get_scene_asset_paths(
+    scene_paths, abs_paths=False, skip_failures=True, progress=None
+):
+    """Extracts all asset paths for the specified 3D scenes.
+
+    Args:
+        scene_paths: an iterable of ``.fo3d`` paths
+        abs_paths (False): whether to return absolute paths
+        skip_failures (True): whether to gracefully continue without raising an
+            error if metadata cannot be computed for a file
+        progress (None): whether to render a progress bar (True/False), use the
+            default value ``fiftyone.config.show_progress_bars`` (None), or a
+            progress callback function to invoke instead
+
+    Returns:
+        a dict mapping scene paths to lists of asset paths
+    """
+    if not scene_paths:
+        return {}
+
+    _scene_paths, all_local = _get_scene_paths(scene_paths)
+
+    if all_local:
+        if progress is None:
+            progress = False
+    else:
+        logger.info("Getting asset paths...")
+
+    _get_scene_asset_paths_single_bound = functools.partial(
+        _get_scene_asset_paths_single,
+        abs_paths=abs_paths,
+        skip_failures=skip_failures,
+    )
+    all_asset_paths = fos.run(
+        _get_scene_asset_paths_single_bound,
+        list(zip(_scene_paths, scene_paths)),
+        progress=progress,
     )
 
-    if pcd_path is None or os.path.isabs(pcd_path):
-        return pcd_path
+    asset_map = dict(zip(scene_paths, all_asset_paths))
 
-    return os.path.join(os.path.dirname(scene_path), pcd_path)
+    return asset_map
 
 
 def compute_orthographic_projection_images(
@@ -569,28 +618,26 @@ def compute_orthographic_projection_images(
 
     fov.validate_collection(view, media_type={fom.POINT_CLOUD, fom.THREE_D})
 
+    if out_group_slice is not None:
+        out_samples = []
+
     filename_maker = fou.UniqueFilenameMaker(
         output_dir=output_dir, rel_dir=rel_dir
     )
 
-    if out_group_slice is not None:
-        out_samples = []
-
     for sample in view.iter_samples(autosave=True, progress=progress):
-        projection_pcd_filepath = sample.filepath
-
         if view.media_type == fom.THREE_D:
-            projection_pcd_filepath = _get_pcd_filepath_from_fo3d_scene(
-                Scene.from_fo3d(sample.filepath), sample.filepath
-            )
+            pcd_filepath = _get_pcd_filepath_from_scene(sample.filepath)
+        else:
+            pcd_filepath = sample.filepath
 
         image_path = filename_maker.get_output_path(
-            projection_pcd_filepath, output_ext=".png"
+            pcd_filepath, output_ext=".png"
         )
 
         try:
             img, metadata = compute_orthographic_projection_image(
-                projection_pcd_filepath,
+                pcd_filepath,
                 size,
                 shading_mode=shading_mode,
                 colormap=colormap,
@@ -738,6 +785,42 @@ def compute_orthographic_projection_image(
     return image, metadata
 
 
+def _get_pcd_filepath_from_scene(scene_path: str):
+    scene = Scene.from_fo3d(scene_path)
+
+    explicitly_flagged_pcd_path = None
+    fallover_pcd_path = None
+
+    def _visit_node_dfs(node):
+        nonlocal explicitly_flagged_pcd_path
+        nonlocal fallover_pcd_path
+
+        if hasattr(node, "pcd_path") and node.flag_for_projection:
+            explicitly_flagged_pcd_path = node.pcd_path
+        else:
+            if hasattr(node, "pcd_path"):
+                fallover_pcd_path = node.pcd_path
+
+            for child in node.children:
+                _visit_node_dfs(child)
+
+    _visit_node_dfs(scene)
+
+    pcd_path = (
+        explicitly_flagged_pcd_path
+        if explicitly_flagged_pcd_path
+        else fallover_pcd_path
+    )
+
+    if pcd_path is None:
+        return None
+
+    if not fos.isabs(pcd_path):
+        pcd_path = fos.join(os.path.dirname(scene_path), pcd_path)
+
+    return fos.resolve(pcd_path)
+
+
 def _parse_point_cloud(
     filepath,
     size=None,

tests/unittests/metadata_tests.py

@@ -0,0 +1,53 @@
+"""
+FiftyOne metadata unit tests.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+import json
+import os
+import tempfile
+import unittest
+
+import fiftyone.core.metadata as fom
+import fiftyone.core.threed as fo3d
+
+
+class SceneMetadataTests(unittest.TestCase):
+    def test_build_for(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            scene_path = os.path.join(temp_dir, "scene.fo3d")
+            files = [("stl", 123), ("obj", 321), ("jpeg", 5151), ("mtl", 867)]
+            for file, num_bytes in files:
+                with open(
+                    os.path.join(temp_dir, f"{file}.{file}"), "wb"
+                ) as of:
+                    of.write(b"A" * num_bytes)
+
+            scene = fo3d.Scene()
+            scene.background = fo3d.SceneBackground(image="jpeg.jpeg")
+            scene.add(fo3d.ObjMesh("blah-obj", "obj.obj", "mtl.mtl"))
+            scene.add(fo3d.StlMesh("blah-stl", "stl.stl"))
+
+            # Add same file again - should not add to size_bytes though,
+            #   even though this is abs and the other was relative
+            scene.add(
+                fo3d.ObjMesh("blah-obj2", os.path.join(temp_dir, "obj.obj"))
+            )
+
+            scene.write(scene_path)
+
+            metadata = fom.SceneMetadata.build_for(scene_path)
+
+            self.assertEqual(metadata.mime_type, "application/octet-stream")
+
+            # Read the scene back again so we'll compare more accurately
+            expected_size = len(
+                json.dumps(fo3d.Scene.from_fo3d(scene_path).as_dict())
+            ) + sum(t[1] for t in files)
+            self.assertEqual(metadata.size_bytes, expected_size)
+            self.assertDictEqual(
+                metadata.asset_counts,
+                {"obj": 2, "jpeg": 1, "stl": 1, "mtl": 1},
+            )