fix: introduce vertices and faces doc

Signed-off-by: anna-charlotte <[email protected]>
docarray · samsja · Feb 17, 2023 · Feb 9, 2023 · Feb 9, 2023 · Feb 9, 2023
commit 0244816584fbddf2a207e5f376d3ec0675014fbd
diff --git a/docarray/documents/mesh/__init__.py b/docarray/documents/mesh/__init__.py
@@ -0,0 +1,3 @@
+from docarray.documents.mesh.mesh import Mesh3D
+
+__all__ = ['Mesh3D']
diff --git a/docarray/documents/mesh.py → docarray/documents/mesh/mesh.py b/docarray/documents/mesh.py → docarray/documents/mesh/mesh.py
@@ -1,7 +1,9 @@
 from typing import Any, Optional, Type, TypeVar, Union
 
 from docarray.base_document import BaseDocument
-from docarray.typing import AnyEmbedding, AnyTensor, Mesh3DUrl
+from docarray.documents.mesh.vertices_and_faces import VerticesAndFaces
+from docarray.typing.tensor.embedding import AnyEmbedding
+from docarray.typing.url.url_3d.mesh_url import Mesh3DUrl
 
 T = TypeVar('T', bound='Mesh3D')
 
@@ -17,9 +19,10 @@ class Mesh3D(BaseDocument):
     tensor of shape (n_faces, 3). Each number in that tensor refers to an index of a
     vertex in the tensor of vertices.
 
-    The Mesh3D Document can contain an Mesh3DUrl (`Mesh3D.url`), an AnyTensor of
-    vertices (`Mesh3D.vertices`), an AnyTensor of faces (`Mesh3D.faces`) and an
-    AnyEmbedding (`Mesh3D.embedding`).
+    The Mesh3D Document can contain an Mesh3DUrl (`Mesh3D.url`), a VerticesAndFaces
+    object containing an AnyTensor of vertices (`Mesh3D.tensors.vertices) and an
+    AnyTensor of faces (`Mesh3D.tensors.faces), and an AnyEmbedding
+    (`Mesh3D.embedding`).
 
     EXAMPLE USAGE:
 
@@ -31,9 +34,9 @@ class Mesh3D(BaseDocument):
 
         # use it directly
         mesh = Mesh3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
-        mesh.vertices, mesh.faces = mesh.url.load()
+        mesh.tensors = mesh.url.load()
         model = MyEmbeddingModel()
-        mesh.embedding = model(mesh.vertices)
+        mesh.embedding = model(mesh.tensors.vertices)
 
     You can extend this Document:
 
@@ -49,7 +52,7 @@ class MyMesh3D(Mesh3D):
 
 
         mesh = MyMesh3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
-        mesh.vertices, mesh.faces = mesh.url.load()
+        mesh.tensors = mesh.url.load()
         model = MyEmbeddingModel()
         mesh.embedding = model(mesh.vertices)
         mesh.name = 'my first mesh'
@@ -72,16 +75,32 @@ class MultiModalDoc(BaseDocument):
             mesh=Mesh3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj'),
             text=Text(text='hello world, how are you doing?'),
         )
-        mmdoc.mesh.vertices, mmdoc.mesh.faces = mmdoc.mesh.url.load()
+        mmdoc.mesh.tensors = mmdoc.mesh.url.load()
 
         # or
         mmdoc.mesh.bytes = mmdoc.mesh.url.load_bytes()
 
+
+    You can display your point cloud from either its url, or its tensors:
+
+    .. code-block:: python
+
+        from docarray.documents import Mesh3D
+
+        # display from url
+        mesh = Mesh3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
+        mesh.url.display()
+
+        # display from tensors
+        mesh.tensors = mesh.url.load()
+        model = MyEmbeddingModel()
+        mesh.embedding = model(mesh.tensors.vertices)
+
+
     """
 
     url: Optional[Mesh3DUrl]
-    vertices: Optional[AnyTensor]
-    faces: Optional[AnyTensor]
+    tensors: Optional[VerticesAndFaces]
     embedding: Optional[AnyEmbedding]
     bytes: Optional[bytes]
 
@@ -93,28 +112,3 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         return super().validate(value)
-
-    def display(self, display_from: str = 'url') -> None:
-        """
-        Plot mesh consisting of vertices and faces.
-        :param display_from: display from either url or tensors (vertices and faces).
-        """
-
-        if display_from not in ['tensor', 'url']:
-            raise ValueError(f'Expected one of ["tensor", "url"], got "{display_from}"')
-
-        if display_from == 'url':
-            if self.url is None:
-                raise ValueError('Can\'t display mesh from url when the url is None.')
-            self.url.display()
-        else:
-            import trimesh
-            from IPython.display import display
-
-            if self.vertices is None or self.faces is None:
-                raise ValueError(
-                    'Can\'t display mesh from tensor when vertices and/or faces is None'
-                )
-
-            mesh = trimesh.Trimesh(vertices=self.vertices, faces=self.faces)
-            display(mesh.show())
diff --git a/docarray/documents/mesh/vertices_and_faces.py b/docarray/documents/mesh/vertices_and_faces.py
@@ -0,0 +1,40 @@
+from typing import Any, Optional, Type, TypeVar, Union
+
+from docarray.base_document import BaseDocument
+from docarray.typing.proto_register import _register_proto
+from docarray.typing.tensor.tensor import AnyTensor
+
+T = TypeVar('T', bound='VerticesAndFaces')
+
+
+@_register_proto(proto_type_name='vertices_and_faces')
+class VerticesAndFaces(BaseDocument):
+    """ """
+
+    vertices: Optional[AnyTensor]
+    faces: Optional[AnyTensor]
+
+    @classmethod
+    def validate(
+        cls: Type[T],
+        value: Union[str, Any],
+    ) -> T:
+        if isinstance(value, str):
+            value = cls(url=value)
+        return super().validate(value)
+
+    def display(self) -> None:
+        """
+        Plot mesh consisting of vertices and faces.
+        """
+        import trimesh
+        from IPython.display import display
+
+        if self.vertices is None or self.faces is None:
+            raise ValueError(
+                'Can\'t display mesh from tensors when the vertices and/or faces '
+                'are None.'
+            )
+
+        mesh = trimesh.Trimesh(vertices=self.vertices, faces=self.faces)
+        display(mesh.show())
diff --git a/docarray/typing/url/url_3d/mesh_url.py b/docarray/typing/url/url_3d/mesh_url.py
@@ -1,4 +1,4 @@
-from typing import NamedTuple, TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
@@ -7,12 +7,10 @@
 from docarray.typing.tensor.ndarray import NdArray
 from docarray.typing.url.url_3d.url_3d import Url3D
 
-T = TypeVar('T', bound='Mesh3DUrl')
-
+if TYPE_CHECKING:
+    from docarray.documents.mesh.vertices_and_faces import VerticesAndFaces
 
-class Mesh3DLoadResult(NamedTuple):
-    vertices: NdArray
-    faces: NdArray
+T = TypeVar('T', bound='Mesh3DUrl')
 
 
 @_register_proto(proto_type_name='mesh_url')
@@ -22,9 +20,9 @@ class Mesh3DUrl(Url3D):
     Can be remote (web) URL, or a local file path.
     """
 
-    def load(self: T) -> Mesh3DLoadResult:
+    def load(self: T) -> 'VerticesAndFaces':
         """
-        Load the data from the url into a named tuple of two NdArrays containing
+        Load the data from the url into a VerticesAndFaces object containing
         vertices and faces information.
 
         EXAMPLE USAGE
@@ -34,7 +32,7 @@ def load(self: T) -> Mesh3DLoadResult:
             from docarray import BaseDocument
             import numpy as np
 
-            from docarray.typing import Mesh3DUrl
+            from docarray.typing import Mesh3DUrl, NdArray
 
 
             class MyDoc(BaseDocument):
@@ -43,19 +41,20 @@ class MyDoc(BaseDocument):
 
             doc = MyDoc(mesh_url="toydata/tetrahedron.obj")
 
-            vertices, faces = doc.mesh_url.load()
-            assert isinstance(vertices, np.ndarray)
-            assert isinstance(faces, np.ndarray)
+            tensors = doc.mesh_url.load()
+            assert isinstance(tensors.vertices, NdArray)
+            assert isinstance(tensors.faces, NdArray)
 
-        :return: named tuple of two NdArrays representing the mesh's vertices and faces
+        :return: VerticesAndFaces object containing vertices and faces information.
         """
+        from docarray.documents.mesh.vertices_and_faces import VerticesAndFaces
 
         mesh = self._load_trimesh_instance(force='mesh')
 
         vertices = parse_obj_as(NdArray, mesh.vertices.view(np.ndarray))
         faces = parse_obj_as(NdArray, mesh.faces.view(np.ndarray))
 
-        return Mesh3DLoadResult(vertices=vertices, faces=faces)
+        return VerticesAndFaces(vertices=vertices, faces=faces)
 
     def display(self) -> None:
         """

diff --git a/tests/integrations/predefined_document/test_mesh.py b/tests/integrations/predefined_document/test_mesh.py
@@ -2,7 +2,7 @@
 import pytest
 from pydantic import parse_obj_as
 
-from docarray import BaseDocument
+from docarray.base_document.document import BaseDocument
 from docarray.documents import Mesh3D
 from tests import TOYDATA_DIR
 
@@ -17,10 +17,10 @@ def test_mesh(file_url):
 
     mesh = Mesh3D(url=file_url)
 
-    mesh.vertices, mesh.faces = mesh.url.load()
+    mesh.tensors = mesh.url.load()
 
-    assert isinstance(mesh.vertices, np.ndarray)
-    assert isinstance(mesh.faces, np.ndarray)
+    assert isinstance(mesh.tensors.vertices, np.ndarray)
+    assert isinstance(mesh.tensors.faces, np.ndarray)
 
 
 def test_str_init():
@@ -37,16 +37,3 @@ class MyDoc(BaseDocument):
 
     assert doc.mesh1.url == 'http://hello.ply'
     assert doc.mesh2.url == 'http://hello.ply'
-
-
-def test_display_illegal_param():
-    mesh = Mesh3D(url='http://myurl.ply')
-    with pytest.raises(ValueError):
-        mesh.display(display_from='tensor')
-
-    mesh = Mesh3D(vertices=np.zeros((10, 3)), faces=np.ones((10, 3)))
-    with pytest.raises(ValueError):
-        mesh.display(display_from='url')
-
-    with pytest.raises(ValueError):
-        mesh.display(display_from='illegal')
diff --git a/tests/units/typing/url/test_mesh_url.py b/tests/units/typing/url/test_mesh_url.py
@@ -4,7 +4,6 @@
 
 from docarray.base_document.io.json import orjson_dumps
 from docarray.typing import Mesh3DUrl, NdArray
-from docarray.typing.url.url_3d.mesh_url import Mesh3DLoadResult
 from tests import TOYDATA_DIR
 
 MESH_FILES = {
@@ -28,14 +27,14 @@
 )
 def test_load(file_format, file_path):
     url = parse_obj_as(Mesh3DUrl, file_path)
-    vertices, faces = url.load()
+    tensors = url.load()
 
-    assert isinstance(vertices, np.ndarray)
-    assert isinstance(vertices, NdArray)
-    assert isinstance(faces, np.ndarray)
-    assert isinstance(faces, NdArray)
-    assert vertices.shape[1] == 3
-    assert faces.shape[1] == 3
+    assert isinstance(tensors.vertices, np.ndarray)
+    assert isinstance(tensors.vertices, NdArray)
+    assert isinstance(tensors.faces, np.ndarray)
+    assert isinstance(tensors.faces, NdArray)
+    assert tensors.vertices.shape[1] == 3
+    assert tensors.faces.shape[1] == 3
 
 
 @pytest.mark.slow
@@ -44,7 +43,7 @@ def test_load(file_format, file_path):
     'file_path',
     [*MESH_FILES.values(), REMOTE_OBJ_FILE],
 )
-@pytest.mark.parametrize('field', [f for f in Mesh3DLoadResult._fields])
+@pytest.mark.parametrize('field', ['vertices', 'faces'])
 def test_load_one_of_fields(file_path, field):
     url = parse_obj_as(Mesh3DUrl, file_path)
     field = getattr(url.load(), field)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from docarray.documents.mesh.mesh import Mesh3D

		__all__ = ['Mesh3D']