docarray · anna-charlotte · Feb 16, 2023 · Feb 9, 2023 · Feb 9, 2023 · Feb 9, 2023
diff --git a/docarray/documents/mesh/__init__.py b/docarray/documents/mesh/__init__.py
@@ -0,0 +1,3 @@
+from docarray.documents.mesh.mesh_3d import Mesh3D
+
+__all__ = ['Mesh3D']
diff --git a/docarray/documents/mesh.py → docarray/documents/mesh/mesh_3d.py b/docarray/documents/mesh.py → docarray/documents/mesh/mesh_3d.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:
 
@@ -43,13 +46,14 @@ class Mesh3D(BaseDocument):
         from docarray.typing import AnyEmbedding
         from typing import Optional
 
+
         # extend it
         class MyMesh3D(Mesh3D):
             name: Optional[Text]
 
 
         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'
@@ -62,6 +66,7 @@ class MyMesh3D(Mesh3D):
         from docarray import BaseDocument
         from docarray.documents import Mesh3D, Text
 
+
         # compose it
         class MultiModalDoc(BaseDocument):
             mesh: Mesh3D
@@ -72,16 +77,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 3D mesh in a notebook 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]
 

diff --git a/docarray/documents/mesh/vertices_and_faces.py b/docarray/documents/mesh/vertices_and_faces.py
@@ -0,0 +1,42 @@
+from typing import Any, Type, TypeVar, Union
+
+from docarray.base_document import BaseDocument
+from docarray.typing.tensor.tensor import AnyTensor
+
+T = TypeVar('T', bound='VerticesAndFaces')
+
+
+class VerticesAndFaces(BaseDocument):
+    """
+    Document for handling 3D mesh tensor data.
+
+    A VerticesAndFaces Document can contain an AnyTensor containing the vertices
+    information (`VerticesAndFaces.vertices`), and an AnyTensor containing the faces
+    information (`VerticesAndFaces.faces`).
+    """
+
+    vertices: AnyTensor
+    faces: AnyTensor
+
+    @classmethod
+    def validate(
+        cls: Type[T],
+        value: Union[str, Any],
+    ) -> T:
+        return super().validate(value)
+
+    def display(self) -> None:
+        """
+        Plot mesh consisting of vertices and faces in notebook.
+        """
+        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/documents/point_cloud/__init__.py b/docarray/documents/point_cloud/__init__.py
@@ -0,0 +1,3 @@
+from docarray.documents.point_cloud.point_cloud_3d import PointCloud3D
+
+__all__ = ['PointCloud3D']
diff --git a/docarray/documents/point_cloud.py → ...y/documents/point_cloud/point_cloud_3d.py b/docarray/documents/point_cloud.py → ...y/documents/point_cloud/point_cloud_3d.py
@@ -3,7 +3,8 @@
 import numpy as np
 
 from docarray.base_document import BaseDocument
-from docarray.typing import AnyEmbedding, AnyTensor, PointCloud3DUrl
+from docarray.documents.point_cloud.points_and_colors import PointsAndColors
+from docarray.typing import AnyEmbedding, PointCloud3DUrl
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.utils.misc import is_tf_available, is_torch_available
 
@@ -27,8 +28,9 @@ class PointCloud3D(BaseDocument):
     representation, the point cloud is a fixed size ndarray (shape=(n_samples, 3)) and
     hence easier for deep learning algorithms to handle.
 
-    A PointCloud3D Document can contain an PointCloud3DUrl (`PointCloud3D.url`), an
-    AnyTensor (`PointCloud3D.tensor`), and an AnyEmbedding (`PointCloud3D.embedding`).
+    A PointCloud3D Document can contain an PointCloud3DUrl (`PointCloud3D.url`),
+    a PointsAndColors object (`PointCloud3D.tensors`), and an AnyEmbedding
+    (`PointCloud3D.embedding`).
 
     EXAMPLE USAGE:
 
@@ -40,9 +42,9 @@ class PointCloud3D(BaseDocument):
 
         # use it directly
         pc = PointCloud3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
-        pc.tensor = pc.url.load(samples=100)
+        pc.tensors = pc.url.load(samples=100)
         model = MyEmbeddingModel()
-        pc.embedding = model(pc.tensor)
+        pc.embedding = model(pc.tensors.points)
 
     You can extend this Document:
 
@@ -58,10 +60,10 @@ class MyPointCloud3D(PointCloud3D):
 
 
         pc = MyPointCloud3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
-        pc.tensor = pc.url.load(samples=100)
+        pc.tensors = pc.url.load(samples=100)
         model = MyEmbeddingModel()
-        pc.embedding = model(pc.tensor)
-        pc.second_embedding = model(pc.tensor)
+        pc.embedding = model(pc.tensors.points)
+        pc.second_embedding = model(pc.tensors.colors)
 
 
     You can use this Document for composition:
@@ -83,16 +85,32 @@ class MultiModalDoc(BaseDocument):
             ),
             text=Text(text='hello world, how are you doing?'),
         )
-        mmdoc.point_cloud.tensor = mmdoc.point_cloud.url.load(samples=100)
+        mmdoc.point_cloud.tensors = mmdoc.point_cloud.url.load(samples=100)
 
         # or
 
         mmdoc.point_cloud.bytes = mmdoc.point_cloud.url.load_bytes()
 
+
+    You can display your point cloud from either its url, or its tensors:
+
+    .. code-block:: python
+
+        from docarray.documents import PointCloud3D
+
+        # display from url
+        pc = PointCloud3D(url='https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj')
+        pc.url.display()
+
+        # display from tensors
+        pc.tensors = pc.url.load(samples=10000)
+        model = MyEmbeddingModel()
+        pc.embedding = model(pc.tensors.points)
+
     """
 
     url: Optional[PointCloud3DUrl]
-    tensor: Optional[AnyTensor]
+    tensors: Optional[PointsAndColors]
     embedding: Optional[AnyEmbedding]
     bytes: Optional[bytes]
 
@@ -108,6 +126,6 @@ def validate(
             and isinstance(value, torch.Tensor)
             or (tf_available and isinstance(value, tf.Tensor))
         ):
-            value = cls(tensor=value)
+            value = cls(tensors=PointsAndColors(points=value))
 
         return super().validate(value)
diff --git a/docarray/documents/point_cloud/points_and_colors.py b/docarray/documents/point_cloud/points_and_colors.py
@@ -0,0 +1,66 @@
+from typing import Any, Optional, Type, TypeVar, Union
+
+import numpy as np
+
+from docarray.base_document import BaseDocument
+from docarray.typing import AnyTensor
+from docarray.typing.tensor.abstract_tensor import AbstractTensor
+from docarray.utils.misc import is_tf_available, is_torch_available
+
+torch_available = is_torch_available()
+if torch_available:
+    import torch
+
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+
+T = TypeVar('T', bound='PointsAndColors')
+
+
+class PointsAndColors(BaseDocument):
+    """
+    Document for handling point clouds tensor data.
+
+    A PointsAndColors Document can contain an AnyTensor containing the points in
+    3D space information (`PointsAndColors.points`), and an AnyTensor containing
+    the points' color information (`PointsAndColors.colors`).
+    """
+
+    points: AnyTensor
+    colors: Optional[AnyTensor]
+
+    @classmethod
+    def validate(
+        cls: Type[T],
+        value: Union[str, AbstractTensor, Any],
+    ) -> T:
+        if isinstance(value, (AbstractTensor, np.ndarray)) or (
+            torch_available
+            and isinstance(value, torch.Tensor)
+            or (tf_available and isinstance(value, tf.Tensor))
+        ):
+            value = cls(points=value)
+
+        return super().validate(value)
+
+    def display(self) -> None:
+        """
+        Plot point cloud consisting of points in 3D space and optionally colors in
+        notebook.
+        """
+        import trimesh
+        from IPython.display import display
+
+        colors = (
+            self.colors
+            if self.colors is not None
+            else np.tile(
+                np.array([0, 0, 0]),
+                (self.points.get_comp_backend().shape(self.points)[0], 1),
+            )
+        )
+        pc = trimesh.points.PointCloud(vertices=self.points, colors=colors)
+
+        s = trimesh.Scene(geometry=pc)
+        display(s.show())
diff --git a/docarray/typing/tensor/tensor.py b/docarray/typing/tensor/tensor.py
@@ -13,11 +13,10 @@
     from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor  # noqa: F401
 
 
+AnyTensor = Union[NdArray]
 if torch_available and tf_available:
-    AnyTensor = Union[NdArray, TorchTensor, TensorFlowTensor]
+    AnyTensor = Union[NdArray, TorchTensor, TensorFlowTensor]  # type: ignore
 elif torch_available:
     AnyTensor = Union[NdArray, TorchTensor]  # type: ignore
 elif tf_available:
     AnyTensor = Union[NdArray, TensorFlowTensor]  # type: ignore
-else:
-    AnyTensor = Union[NdArray]  # type: ignore
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,16 +41,28 @@ 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:
+        """
+        Plot mesh in notebook from url.
+        This loads the Trimesh instance of the 3D mesh, and then displays it.
+        """
+        from IPython.display import display
+
+        mesh = self._load_trimesh_instance()
+        display(mesh.show())
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from docarray.documents.mesh.mesh_3d import Mesh3D

		__all__ = ['Mesh3D']
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from docarray.documents.point_cloud.point_cloud_3d import PointCloud3D

		__all__ = ['PointCloud3D']