From 4f54c643b3f0f41a370139716957e9bf1bd6350d Mon Sep 17 00:00:00 2001 From: numb3r3 Date: Tue, 11 Jan 2022 15:07:30 +0800 Subject: [PATCH] fix: improve mesh data loading --- docarray/document/mixins/mesh.py | 19 +++++++++---------- tests/unit/document/test_converters.py | 4 ++++ 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/docarray/document/mixins/mesh.py b/docarray/document/mixins/mesh.py index e0ca228f8b0..b6ad3a27fab 100644 --- a/docarray/document/mixins/mesh.py +++ b/docarray/document/mixins/mesh.py @@ -22,18 +22,17 @@ def load_uri_to_point_cloud_blob( """ import trimesh - mesh = trimesh.load_mesh(self.uri).deduplicated() - - pcs = [] - for geo in mesh.geometry.values(): - geo: trimesh.Trimesh - pcs.append(geo.sample(samples)) - if as_chunks: from .. import Document - for p in pcs: - self.chunks.append(Document(blob=p)) + # try to coerce everything into a scene + scene = trimesh.load(self.uri, force='scene') + for geo in scene.geometry.values(): + geo: trimesh.Trimesh + self.chunks.append(Document(blob=geo.sample(samples))) else: - self.blob = np.stack(pcs).squeeze() + # combine a scene into a single mesh + mesh = trimesh.load(self.uri, force='mesh') + self.blob = mesh.sample(samples) + return self diff --git a/tests/unit/document/test_converters.py b/tests/unit/document/test_converters.py index 37923f84e1e..a004dda9b59 100644 --- a/tests/unit/document/test_converters.py +++ b/tests/unit/document/test_converters.py @@ -234,3 +234,7 @@ def test_glb_converters(): doc = Document(uri=os.path.join(cur_dir, 'toydata/test.glb')) doc.load_uri_to_point_cloud_blob(2000) assert doc.blob.shape == (2000, 3) + + doc.load_uri_to_point_cloud_blob(2000, as_chunks=True) + assert len(doc.chunks) == 1 + assert doc.chunks[0].blob.shape == (2000, 3)