feat: display mesh and pointcloud

Signed-off-by: anna-charlotte <[email protected]>
docarray · anna-charlotte · Feb 16, 2023 · Feb 9, 2023 · Feb 9, 2023 · Feb 9, 2023
commit 943d63698c8acfc9dd5a076c76893a2e9eb4b250
diff --git a/docarray/documents/mesh.py b/docarray/documents/mesh.py
@@ -93,3 +93,18 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         return super().validate(value)
+
+    def display(self):
+        """Plot mesh consisting of vertices and faces."""
+        from IPython.display import display
+
+        if self.url:
+            # mesh from uri
+            mesh = self.url._load_trimesh_instance()
+            display(mesh.show())
+        else:
+            # mesh from vertices and faces tensors
+            import trimesh
+
+            mesh = trimesh.Trimesh(vertices=self.vertices, faces=self.faces)
+            display(mesh.show())
diff --git a/docarray/documents/point_cloud.py b/docarray/documents/point_cloud.py
@@ -93,6 +93,7 @@ class MultiModalDoc(BaseDocument):
 
     url: Optional[PointCloud3DUrl]
     tensor: Optional[AnyTensor]
+    color_tensor: Optional[AnyTensor]
     embedding: Optional[AnyEmbedding]
     bytes: Optional[bytes]
 
@@ -111,3 +112,26 @@ def validate(
             value = cls(tensor=value)
 
         return super().validate(value)
+
+    def display(self) -> None:
+        """Plot interactive point cloud from :attr:`.tensor`"""
+        import trimesh
+        from hubble.utils.notebook import is_notebook
+        from IPython.display import display
+
+        colors = (
+            self.color_tensor
+            if self.color_tensor
+            else np.tile(np.array([0, 0, 0]), (len(self.tensor), 1))
+        )
+
+        pc = trimesh.points.PointCloud(
+            vertices=self.tensor,
+            colors=colors,
+        )
+
+        if is_notebook():
+            s = trimesh.Scene(geometry=pc)
+            display(s.show())
+        else:
+            display(pc.show())