docarray · JoanFM · Nov 9, 2022 · Nov 2, 2022 · Nov 2, 2022 · Nov 2, 2022
diff --git a/docarray/document/mixins/mesh.py b/docarray/document/mixins/mesh.py
@@ -1,15 +1,44 @@
-import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
 
 if TYPE_CHECKING:  # pragma: no cover
     from docarray.typing import T
+    import trimesh
+
+
+class Mesh:
+    FILE_EXTENSIONS = [
+        'glb',
+        'obj',
+        'ply',
+    ]
+    VERTICES = 'vertices'
+    FACES = 'faces'
 
 
 class MeshDataMixin:
     """Provide helper functions for :class:`Document` to support 3D mesh data and point cloud."""
 
+    def _load_mesh(
+        self, force: str = None
+    ) -> Union['trimesh.Trimesh', 'trimesh.Scene']:
+        """Load a trimesh.Mesh or trimesh.Scene object from :attr:`.uri`.
+
+        :param force: str or None. For 'mesh' try to coerce scenes into a single mesh. For 'scene'
+            try to coerce everything into a scene.
+        :return: trimesh.Mesh or trimesh.Scene object
+        """
+        import urllib.parse
+        import trimesh
+
+        scheme = urllib.parse.urlparse(self.uri).scheme
+        loader = trimesh.load_remote if scheme in ['http', 'https'] else trimesh.load
+
+        mesh = loader(self.uri, force=force)
+
+        return mesh
+
     def load_uri_to_point_cloud_tensor(
         self: 'T', samples: int, as_chunks: bool = False
     ) -> 'T':
@@ -21,23 +50,19 @@ def load_uri_to_point_cloud_tensor(
 
         :return: itself after processed
         """
-        import trimesh
-        import urllib.parse
-
-        scheme = urllib.parse.urlparse(self.uri).scheme
-        loader = trimesh.load_remote if scheme in ['http', 'https'] else trimesh.load
 
         if as_chunks:
+            import trimesh
             from docarray.document import Document
 
             # try to coerce everything into a scene
-            scene = loader(self.uri, force='scene')
+            scene = self._load_mesh(force='scene')
             for geo in scene.geometry.values():
                 geo: trimesh.Trimesh
                 self.chunks.append(Document(tensor=np.array(geo.sample(samples))))
         else:
             # combine a scene into a single mesh
-            mesh = loader(self.uri, force='mesh')
+            mesh = self._load_mesh(force='mesh')
             self.tensor = np.array(mesh.sample(samples))
 
         return self
@@ -47,22 +72,16 @@ def load_uri_to_vertices_and_faces(self: 'T') -> 'T':
 
         :return: itself after processed
         """
-
-        import trimesh
-        import urllib.parse
         from docarray.document import Document
 
-        scheme = urllib.parse.urlparse(self.uri).scheme
-        loader = trimesh.load_remote if scheme in ['http', 'https'] else trimesh.load
-
-        mesh = loader(self.uri, force='mesh')
+        mesh = self._load_mesh(force='mesh')
 
         vertices = mesh.vertices.view(np.ndarray)
         faces = mesh.faces.view(np.ndarray)
 
         self.chunks = [
-            Document(name='vertices', tensor=vertices),
-            Document(name='faces', tensor=faces),
+            Document(name=Mesh.VERTICES, tensor=vertices),
+            Document(name=Mesh.FACES, tensor=faces),
         ]
 
         return self
@@ -73,18 +92,18 @@ def load_vertices_and_faces_to_point_cloud(self: 'T', samples: int) -> 'T':
         :param samples: number of points to sample from the mesh
         :return: itself after processed
         """
-        import trimesh
-
         vertices = None
         faces = None
 
         for chunk in self.chunks:
-            if chunk.tags['name'] == 'vertices':
+            if chunk.tags['name'] == Mesh.VERTICES:
                 vertices = chunk.tensor
-            if chunk.tags['name'] == 'faces':
+            if chunk.tags['name'] == Mesh.FACES:
                 faces = chunk.tensor
 
         if vertices is not None and faces is not None:
+            import trimesh
+
             mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
             self.tensor = np.array(mesh.sample(samples))
         else:

diff --git a/docarray/document/mixins/plot.py b/docarray/document/mixins/plot.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from docarray.helper import deprecate_by
+from docarray.document.mixins.mesh import Mesh
 
 
 class PlotMixin:
@@ -73,26 +73,81 @@ def _plot_recursion(self, tree=None):
     def display(self, from_: Optional[str] = None):
         """
         Plot image data from :attr:`.uri` or from :attr:`.tensor` if :attr:`.uri` is empty .
-
         :param from_: an optional string to decide if a document should display using either the uri or the tensor field.
         """
+        if self._is_3d():
+            self.display_3d()
+        else:
+            if not from_:
+                if self.uri:
+                    from_ = 'uri'
+                elif self.tensor is not None:
+                    from_ = 'tensor'
+                else:
+                    self.summary()
 
-        if not from_:
-            if self.uri:
-                from_ = 'uri'
-            elif self.tensor is not None:
-                from_ = 'tensor'
+            if from_ == 'uri':
+                self.display_uri()
+            elif from_ == 'tensor':
+                self.display_tensor()
             else:
                 self.summary()
 
-        if from_ == 'uri':
-            self.display_uri()
-        elif from_ == 'tensor':
-            self.display_tensor()
+    def _is_3d(self) -> bool:
+        """
+        Tells if Document stores a 3D object saved as point cloud or vertices and face.
+        :return: bool.
+        """
+        if self.uri and self.uri.endswith(tuple(Mesh.FILE_EXTENSIONS)):
+            return True
+        elif (
+            self.tensor is not None
+            and self.tensor.shape[1] == 3
+            and self.tensor.ndim == 2
+        ):
+            return True
+        elif self.chunks is not None:
+            name_tags = [c.tags['name'] for c in self.chunks]
+            if Mesh.VERTICES in name_tags and Mesh.FACES in name_tags:
+                return True
         else:
-            self.summary()
+            return False
 
-    def display_tensor(self):
+    def display_3d(self) -> None:
+        """Plot 3d data."""
+        from IPython.display import display
+        import trimesh
+
+        if self.tensor is not None:
+            # point cloud from tensor
+            from hubble.utils.notebook import is_notebook
+
+            if is_notebook():
+                pc = trimesh.points.PointCloud(
+                    vertices=self.tensor,
+                    colors=np.tile(np.array([0, 0, 0, 1]), (len(self.tensor), 1)),
+                )
+                s = trimesh.Scene(geometry=pc)
+                display(s.show())
+            else:
+                pc = trimesh.points.PointCloud(vertices=self.tensor)
+                display(pc.show())
+
+        elif self.uri:
+            # mesh from uri
+            mesh = self._load_mesh()
+            display(mesh.show())
+
+        elif self.chunks is not None:
+            # mesh from chunks
+            vertices = [
+                c.tensor for c in self.chunks if c.tags['name'] == Mesh.VERTICES
+            ][-1]
+            faces = [c.tensor for c in self.chunks if c.tags['name'] == Mesh.FACES][-1]
+            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+            display(mesh.show())
+
+    def display_tensor(self) -> None:
         """Plot image data from :attr:`.tensor`"""
         if self.tensor is None:
             raise ValueError(