diff --git a/docarray/document/mixins/mesh.py b/docarray/document/mixins/mesh.py index 8f5de66918b..6997f3a192c 100644 --- a/docarray/document/mixins/mesh.py +++ b/docarray/document/mixins/mesh.py @@ -85,8 +85,8 @@ def load_uri_to_vertices_and_faces(self: 'T') -> 'T': faces = mesh.faces.view(np.ndarray) self.chunks = [ - Document(name=MeshEnum.VERTICES, tensor=vertices), - Document(name=MeshEnum.FACES, tensor=faces), + Document(name=MeshEnum.VERTICES.value, tensor=vertices), + Document(name=MeshEnum.FACES.value, tensor=faces), ] return self @@ -101,9 +101,9 @@ def load_vertices_and_faces_to_point_cloud(self: 'T', samples: int) -> 'T': faces = None for chunk in self.chunks: - if chunk.tags['name'] == MeshEnum.VERTICES: + if chunk.tags['name'] == MeshEnum.VERTICES.value: vertices = chunk.tensor - if chunk.tags['name'] == MeshEnum.FACES: + if chunk.tags['name'] == MeshEnum.FACES.value: faces = chunk.tensor if vertices is not None and faces is not None: diff --git a/docarray/document/mixins/plot.py b/docarray/document/mixins/plot.py index 00f6570c6e0..f9c6f2e52b6 100644 --- a/docarray/document/mixins/plot.py +++ b/docarray/document/mixins/plot.py @@ -133,7 +133,10 @@ def _is_3d_vertices_and_faces(self): """ if self.chunks is not None: name_tags = [c.tags['name'] for c in self.chunks] - if MeshEnum.VERTICES in name_tags and MeshEnum.FACES in name_tags: + if ( + MeshEnum.VERTICES.value in name_tags + and MeshEnum.FACES.value in name_tags + ): return True else: return False @@ -173,11 +176,13 @@ def display_vertices_and_faces(self): import trimesh vertices = [ - c.tensor for c in self.chunks if c.tags['name'] == MeshEnum.VERTICES + c.tensor + for c in self.chunks + if c.tags['name'] == MeshEnum.VERTICES.value + ][-1] + faces = [ + c.tensor for c in self.chunks if c.tags['name'] == MeshEnum.FACES.value ][-1] - faces = [c.tensor for c in self.chunks if c.tags['name'] == MeshEnum.FACES][ - -1 - ] mesh = trimesh.Trimesh(vertices=vertices, faces=faces) display(mesh.show()) @@ -191,7 +196,7 @@ def display_point_cloud_tensor(self) -> None: for chunk in self.chunks: if ( 'name' in chunk.tags.keys() - and chunk.tags['name'] == PointCloudEnum.COLORS + and chunk.tags['name'] == PointCloudEnum.COLORS.value and chunk.tensor.shape[-1] in [3, 4] ): colors = chunk.tensor diff --git a/tests/unit/document/test_converters.py b/tests/unit/document/test_converters.py index 9316c4a9b6f..28d8922b704 100644 --- a/tests/unit/document/test_converters.py +++ b/tests/unit/document/test_converters.py @@ -6,7 +6,7 @@ from docarray import Document from docarray.document.generators import from_files -from docarray.document.mixins.mesh import MeshEnum +from docarray.document.mixins.mesh import MeshEnum, PointCloudEnum __windows__ = sys.platform == 'win32' @@ -321,9 +321,9 @@ def test_load_uri_to_vertices_and_faces(uri): doc.load_uri_to_vertices_and_faces() assert len(doc.chunks) == 2 - assert doc.chunks[0].tags['name'] == MeshEnum.VERTICES + assert doc.chunks[0].tags['name'] == MeshEnum.VERTICES.value assert doc.chunks[0].tensor.shape[1] == 3 - assert doc.chunks[1].tags['name'] == MeshEnum.FACES + assert doc.chunks[1].tags['name'] == MeshEnum.FACES.value assert doc.chunks[1].tensor.shape[1] == 3