Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit 227b2e5

Browse files
anna-charlotteanna-charlottealexcg1
authored
feat: add rgbd representation of 3d mesh (#753)
* test: add test and toydata for load uri to rgbd chunk tensors Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * feat: load chunk uris to rgbd tensor Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * fix: change assertions to value errors Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * test: raise exception when uri attr is empty string Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * refactor: to load uris to rgbd tensor Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * feat: add display rgbd image and extract display point cloud Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * style: black formatting Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> * docs: add documentation for rgbd representation Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: change width of rgdb display image Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: change width of rgbd display image Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: adjust size of rgbd display image Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * refactor: extract individual display functions for each 3d repr Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: add documentation of rgbd notebook support Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * fix: use np concatenate instead of np append Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * fix: add epsilon in case max of depth is zero Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: update suggestion Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: apply suggestions Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: apply suggestions Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: apply suggestions Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: apply suggestions Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * test: fix typo in test Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * feat: add legend to depth image Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: update rgbd images with legend Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> Signed-off-by: anna-charlotte <anna-charlotte@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> Co-authored-by: anna-charlotte <anna-charlotte@users.noreply.github.com> Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
1 parent 8470a0d commit 227b2e5

9 files changed

Lines changed: 224 additions & 46 deletions

File tree

‎docarray/document/mixins/mesh.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,3 +112,34 @@ def load_vertices_and_faces_to_point_cloud(self: 'T', samples: int) -> 'T':
112112
)
113113

114114
return self
115+
116+
def load_uris_to_rgbd_tensor(self: 'T') -> 'T':
117+
"""Load RGB image from :attr:`.uri` of :attr:`.chunks[0]` and depth image from :attr:`.uri` of :attr:`.chunks[1]` and merge them into :attr:`.tensor`.
118+
119+
:return: itself after processed
120+
"""
121+
from PIL import Image
122+
123+
if len(self.chunks) != 2:
124+
raise ValueError(
125+
f'The provided Document does not have two chunks but instead {len(self.chunks)}. To load uris to RGBD tensor, the Document needs to have two chunks, with the first one providing the RGB image uri, and the second one providing the depth image uri.'
126+
)
127+
for chunk in self.chunks:
128+
if chunk.uri == '':
129+
raise ValueError(
130+
'A chunk of the given Document does not provide a uri.'
131+
)
132+
133+
rgb_img = np.array(Image.open(self.chunks[0].uri).convert('RGB'))
134+
depth_img = np.array(Image.open(self.chunks[1].uri))
135+
136+
if rgb_img.shape[0:2] != depth_img.shape:
137+
raise ValueError(
138+
f'The provided RGB image and depth image are not of the same shapes: {rgb_img.shape[0:2]} != {depth_img.shape}'
139+
)
140+
141+
self.tensor = np.concatenate(
142+
(rgb_img, np.expand_dims(depth_img, axis=2)), axis=-1
143+
)
144+
145+
return self

‎docarray/document/mixins/plot.py‎

Lines changed: 95 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,12 @@ def display(self, from_: Optional[str] = None):
7676
Plot image data from :attr:`.uri` or from :attr:`.tensor` if :attr:`.uri` is empty .
7777
:param from_: an optional string to decide if a document should display using either the uri or the tensor field.
7878
"""
79-
if self._is_3d():
80-
self.display_3d()
79+
if self._is_3d_point_cloud():
80+
self.display_point_cloud_tensor()
81+
elif self._is_3d_rgbd():
82+
self.display_rgbd_tensor()
83+
elif self._is_3d_vertices_and_faces():
84+
self.display_vertices_and_faces()
8185
else:
8286
if not from_:
8387
if self.uri:
@@ -94,60 +98,46 @@ def display(self, from_: Optional[str] = None):
9498
else:
9599
self.summary()
96100

97-
def _is_3d(self) -> bool:
101+
def _is_3d_point_cloud(self):
98102
"""
99-
Tells if Document stores a 3D object saved as point cloud or vertices and face.
103+
Tells if Document stores a 3D object saved as point cloud tensor.
100104
:return: bool.
101105
"""
102-
if self.uri and self.uri.endswith(tuple(Mesh.FILE_EXTENSIONS)):
103-
return True
104-
elif (
106+
if (
105107
self.tensor is not None
106-
and self.tensor.shape[1] == 3
107108
and self.tensor.ndim == 2
109+
and self.tensor.shape[-1] == 3
110+
):
111+
return True
112+
else:
113+
return False
114+
115+
def _is_3d_rgbd(self):
116+
"""
117+
Tells if Document stores a 3D object saved as RGB-D image tensor.
118+
:return: bool.
119+
"""
120+
if (
121+
self.tensor is not None
122+
and self.tensor.ndim == 3
123+
and self.tensor.shape[-1] == 4
108124
):
109125
return True
110-
elif self.chunks is not None:
126+
else:
127+
return False
128+
129+
def _is_3d_vertices_and_faces(self):
130+
"""
131+
Tells if Document stores a 3D object saved as vertices and faces.
132+
:return: bool.
133+
"""
134+
if self.chunks is not None:
111135
name_tags = [c.tags['name'] for c in self.chunks]
112136
if Mesh.VERTICES in name_tags and Mesh.FACES in name_tags:
113137
return True
114138
else:
115139
return False
116140

117-
def display_3d(self) -> None:
118-
"""Plot 3d data."""
119-
from IPython.display import display
120-
import trimesh
121-
122-
if self.tensor is not None:
123-
# point cloud from tensor
124-
from hubble.utils.notebook import is_notebook
125-
126-
if is_notebook():
127-
pc = trimesh.points.PointCloud(
128-
vertices=self.tensor,
129-
colors=np.tile(np.array([0, 0, 0, 1]), (len(self.tensor), 1)),
130-
)
131-
s = trimesh.Scene(geometry=pc)
132-
display(s.show())
133-
else:
134-
pc = trimesh.points.PointCloud(vertices=self.tensor)
135-
display(pc.show())
136-
137-
elif self.uri:
138-
# mesh from uri
139-
mesh = self._load_mesh()
140-
display(mesh.show())
141-
142-
elif self.chunks is not None:
143-
# mesh from chunks
144-
vertices = [
145-
c.tensor for c in self.chunks if c.tags['name'] == Mesh.VERTICES
146-
][-1]
147-
faces = [c.tensor for c in self.chunks if c.tags['name'] == Mesh.FACES][-1]
148-
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
149-
display(mesh.show())
150-
151141
def display_tensor(self) -> None:
152142
"""Plot image data from :attr:`.tensor`"""
153143
if self.tensor is None:
@@ -169,6 +159,68 @@ def display_tensor(self) -> None:
169159

170160
plt.matshow(self.tensor)
171161

162+
def display_vertices_and_faces(self):
163+
"""Plot mesh consisting of vertices and faces."""
164+
from IPython.display import display
165+
166+
if self.uri:
167+
# mesh from uri
168+
mesh = self._load_mesh()
169+
display(mesh.show())
170+
171+
else:
172+
# mesh from chunks
173+
import trimesh
174+
175+
vertices = [
176+
c.tensor for c in self.chunks if c.tags['name'] == Mesh.VERTICES
177+
][-1]
178+
faces = [c.tensor for c in self.chunks if c.tags['name'] == Mesh.FACES][-1]
179+
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
180+
display(mesh.show())
181+
182+
def display_point_cloud_tensor(self) -> None:
183+
"""Plot interactive point cloud from :attr:`.tensor`"""
184+
import trimesh
185+
from IPython.display import display
186+
from hubble.utils.notebook import is_notebook
187+
188+
if is_notebook():
189+
pc = trimesh.points.PointCloud(
190+
vertices=self.tensor,
191+
colors=np.tile(np.array([0, 0, 0, 1]), (len(self.tensor), 1)),
192+
)
193+
s = trimesh.Scene(geometry=pc)
194+
display(s.show())
195+
else:
196+
pc = trimesh.points.PointCloud(vertices=self.tensor)
197+
display(pc.show())
198+
199+
def display_rgbd_tensor(self) -> None:
200+
"""Plot an RGB-D image and a corresponding depth image from :attr:`.tensor`"""
201+
import matplotlib.pyplot as plt
202+
from mpl_toolkits.axes_grid1 import make_axes_locatable
203+
204+
rgb_img = self.tensor[:, :, :3]
205+
206+
depth_img = self.tensor[:, :, -1]
207+
depth_img = depth_img / (np.max(depth_img) + 1e-08) * 255
208+
depth_img = depth_img.astype(np.uint8)
209+
210+
f, ax = plt.subplots(1, 2, figsize=(16, 6))
211+
212+
ax[0].imshow(rgb_img, interpolation='None')
213+
ax[0].set_title('RGB image\n', fontsize=16)
214+
215+
im2 = ax[1].imshow(self.tensor[:, :, -1], cmap='gray')
216+
cax = make_axes_locatable(ax[1]).append_axes('right', size='5%', pad=0.05)
217+
f.colorbar(im2, cax=cax, orientation='vertical', label='Depth')
218+
219+
ax[1].imshow(depth_img, cmap='gray')
220+
ax[1].set_title('Depth image\n', fontsize=16)
221+
222+
plt.show()
223+
172224
def display_uri(self):
173225
"""Plot image data from :attr:`.uri`"""
174226

‎docs/datatypes/mesh/index.md‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ chunk.tags = {'name': 'faces'}
5151
```
5252

5353

54-
You can display your 3d object and interact with it via:
54+
You can display your 3D object and interact with it via:
5555
```python
5656
doc.display()
5757
```
@@ -1321,7 +1321,7 @@ print(doc.tensor.shape)
13211321
(1000, 3)
13221322
```
13231323

1324-
You can display your 3d object and interact with it via:
1324+
You can display your 3D object and interact with it via:
13251325

13261326
```python
13271327
doc.display()
@@ -2574,3 +2574,35 @@ function animate(){requestAnimationFrame(animate);controls.update();}
25742574
function render(){tracklight.position.copy(camera.position);renderer.render(scene,camera);}
25752575
init();</script></body>
25762576
</html>" width="100%" height="500px" style="border:none;"></iframe
2577+
2578+
2579+
## RGB-D image representation
2580+
2581+
The RGB-D image representation includes an RGB image of shape (w, h, 3) and a corresponding depth image (w, h). The depth image describes the distance between the image plane and the corresponding object for each pixel in the RGB image. Since the RGB and depth image are of identical width and height, they can be easily concatenated and stored in a tensor of shape (w, h, 4). Due to their fixed size, RGB-D images are suitable for 3D data representations for input to machine learning models.
2582+
2583+
With DocArray you can store the uris of an RGB image and its corresponding depth image to the `.uri` attribute of a Document's `.chunks`. You can then load the uris to the Document's `.tensor` attribute at top-level:
2584+
2585+
```python
2586+
from docarray import Document
2587+
2588+
doc = Document(chunks=[Document(uri='rgb_000.jpg'), Document(uri='depth_000.jpg')])
2589+
doc.load_uris_to_rgbd_tensor()
2590+
2591+
doc.summary()
2592+
```
2593+
2594+
```text
2595+
<Document ('id', 'chunks', 'tensor') at 7f907d786d6c11ec840a1e008a366d49>
2596+
└─ chunks
2597+
├─ <Document ('id', 'parent_id', 'granularity', 'uri') at 7f907ab26d6c11ec840a1e008a366d49>
2598+
└─ <Document ('id', 'parent_id', 'granularity', 'uri') at 7f907c106d6c11ec840a1e008a366d49>
2599+
```
2600+
2601+
To display the RGB image and its corresponding depth image:
2602+
2603+
```python
2604+
doc.display()
2605+
```
2606+
2607+
```{figure} rgbd_chair.png
2608+
```

‎docs/datatypes/mesh/rgbd_chair.png‎

1.51 MB
Loading

‎docs/fundamentals/notebook-support/image-rgbd.png‎

1010 KB
Loading

‎docs/fundamentals/notebook-support/index.md‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,16 @@ Video and audio Document can be displayed as well, you can play them in the cell
4141
```{figure} audio-video.png
4242
```
4343

44-
You can also display your 3d object and interact with it in the cell, whether you stored it as a point cloud or vertices and faces.
44+
You can also display your 3D object and interact with it in the cell, whether you stored it as a point cloud or vertices and faces.
4545

4646
```{figure} mesh-point-cloud.png
4747
```
4848

49+
Additionally, you can also display an RGB image and its corresponding depth image:
50+
51+
```{figure} image-rgbd.png
52+
```
53+
4954
## Display DocumentArray
5055

5156
A cell with a DocumentArray object can be pretty-printed automatically.

‎tests/unit/document/test_converters.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,3 +330,61 @@ def test_load_to_point_cloud_without_vertices_faces_set_raise_warning(uri):
330330
AttributeError, match='vertices and faces chunk tensor have not been set'
331331
):
332332
doc.load_vertices_and_faces_to_point_cloud(100)
333+
334+
335+
@pytest.mark.parametrize(
336+
'uri_rgb, uri_depth',
337+
[
338+
(
339+
os.path.join(cur_dir, 'toydata/test_rgb.jpg'),
340+
os.path.join(cur_dir, 'toydata/test_depth.png'),
341+
)
342+
],
343+
)
344+
def test_load_uris_to_rgbd_tensor(uri_rgb, uri_depth):
345+
doc = Document(
346+
chunks=[
347+
Document(uri=uri_rgb),
348+
Document(uri=uri_depth),
349+
]
350+
)
351+
doc.load_uris_to_rgbd_tensor()
352+
353+
assert doc.tensor.shape[-1] == 4
354+
355+
356+
@pytest.mark.parametrize(
357+
'uri_rgb, uri_depth',
358+
[
359+
(
360+
os.path.join(cur_dir, 'toydata/test.png'),
361+
os.path.join(cur_dir, 'toydata/test_depth.png'),
362+
)
363+
],
364+
)
365+
def test_load_uris_to_rgbd_tensor_different_shapes_raise_exception(uri_rgb, uri_depth):
366+
doc = Document(
367+
chunks=[
368+
Document(uri=uri_rgb),
369+
Document(uri=uri_depth),
370+
]
371+
)
372+
with pytest.raises(
373+
ValueError,
374+
match='The provided RGB image and depth image are not of the same shapes',
375+
):
376+
doc.load_uris_to_rgbd_tensor()
377+
378+
379+
def test_load_uris_to_rgbd_tensor_doc_wo_uri_raise_exception():
380+
doc = Document(
381+
chunks=[
382+
Document(),
383+
Document(),
384+
]
385+
)
386+
with pytest.raises(
387+
ValueError,
388+
match='A chunk of the given Document does not provide a uri.',
389+
):
390+
doc.load_uris_to_rgbd_tensor()

‎tests/unit/document/toydata/test_depth.png‎

228 KB
Loading

‎tests/unit/document/toydata/test_rgb.jpg‎

58.9 KB
Loading

0 commit comments

Comments
 (0)