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

Commit 8ceff1e

Browse files
anna-charlottealexcg1
andauthored
feat: add support to load multi page tiff files into chunks (#845)
* test: add test and toydata for loading multi page tiff file Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * feat: store images from multi page tif file in chunks Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * fix: add move channel to axis Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * fix: get n frame in try block Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: update documentation Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> * docs: apply suggestions from code review Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai> Co-authored-by: Alex Cureton-Griffiths <alexcg1@users.noreply.github.com>
1 parent 523dd61 commit 8ceff1e

4 files changed

Lines changed: 86 additions & 10 deletions

File tree

docarray/document/mixins/image.py

Lines changed: 51 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import math
44
import struct
55
import warnings
6-
from typing import Optional, Tuple, Union, BinaryIO, TYPE_CHECKING
6+
from typing import Optional, Tuple, Union, BinaryIO, TYPE_CHECKING, List
77

88
import numpy as np
99

@@ -186,10 +186,25 @@ def load_uri_to_image_tensor(
186186
187187
:return: itself after processed
188188
"""
189+
from docarray import Document, DocumentArray
189190

190191
buffer = _uri_to_blob(self.uri, **kwargs)
191192
tensor = _to_image_tensor(io.BytesIO(buffer), width=width, height=height)
192-
self.tensor = _move_channel_axis(tensor, original_channel_axis=channel_axis)
193+
194+
if isinstance(tensor, np.ndarray):
195+
self.tensor = _move_channel_axis(tensor, original_channel_axis=channel_axis)
196+
elif isinstance(tensor, List):
197+
self.chunks = DocumentArray(
198+
[
199+
Document(
200+
tensor=_move_channel_axis(
201+
tensor[i], original_channel_axis=channel_axis
202+
)
203+
)
204+
for i in range(len(tensor))
205+
]
206+
)
207+
193208
return self
194209

195210
def set_image_tensor_inv_normalization(
@@ -359,26 +374,53 @@ def _to_image_tensor(
359374
source,
360375
width: Optional[int] = None,
361376
height: Optional[int] = None,
362-
) -> 'np.ndarray':
377+
) -> Union[np.ndarray, List[np.array]]:
363378
"""
364-
Convert an image blob to tensor
379+
Convert an image blob to tensor or List of image tensors, if multiple images are stored in file.
365380
366381
:param source: binary blob or file path
367382
:param width: the width of the image tensor.
368383
:param height: the height of the tensor.
369-
:return: image tensor
384+
:return: image tensor or list of image tensors
370385
"""
371386
from PIL import Image
372387

373388
raw_img = Image.open(source)
389+
374390
if width or height:
375391
new_width = width or raw_img.width
376392
new_height = height or raw_img.height
377-
raw_img = raw_img.resize((new_width, new_height))
393+
394+
# support multi page tiff images
378395
try:
379-
return np.array(raw_img.convert('RGB'))
380-
except:
381-
return np.array(raw_img)
396+
n_frames = raw_img.n_frames
397+
except AttributeError:
398+
n_frames = 1
399+
400+
if n_frames > 1:
401+
402+
frames = []
403+
for i in range(raw_img.n_frames):
404+
raw_img.seek(i)
405+
try:
406+
img = raw_img.convert('RGB')
407+
except:
408+
img = raw_img
409+
410+
if width or height:
411+
frames.append(np.array(img.resize((new_width, new_height))))
412+
else:
413+
frames.append(np.array(img))
414+
415+
return frames
416+
417+
else:
418+
if width or height:
419+
raw_img = raw_img.resize((new_width, new_height))
420+
try:
421+
return np.array(raw_img.convert('RGB'))
422+
except:
423+
return np.array(raw_img)
382424

383425

384426
def _to_image_buffer(arr: 'np.ndarray', image_format: str) -> bytes:

docs/datatypes/image/index.md

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,26 @@ print(d.tensor, d.tensor.shape)
3333
(618, 641, 3)
3434
```
3535

36+
DocArray also supports loading multi-page tiff files. In this case, the image tensors are stored to the `.tensor` attributes at the chunk-level instead of the top-level.
37+
38+
```python
39+
from docarray import Document
40+
41+
d = Document(uri='muti_page_tiff_file.tiff')
42+
d.load_uri_to_image_tensor()
43+
44+
d.summary()
45+
```
46+
47+
```text
48+
<Document ('id', 'uri', 'chunks') at 7f907d786d6c11ec840a1e008a366d49>
49+
└─ chunks
50+
├─ <Document ('id', 'parent_id', 'granularity', 'tensor') at 7aa4c0ba66cf6c300b7f07fdcbc2fdc8>
51+
├─ <Document ('id', 'parent_id', 'granularity', 'tensor') at bc94a3e3ca60352f2e4c9ab1b1bb9c22>
52+
└─ <Document ('id', 'parent_id', 'granularity', 'tensor') at 36fe0d1daf4442ad6461c619f8bb25b7>
53+
```
54+
55+
3656
## Simple image processing
3757

3858
DocArray provides some functions to help you preprocess the image data. You can resize it (i.e. downsampling/upsampling) and normalize it; you can switch the channel axis of the `.tensor` to meet certain requirements of other framework; and finally you can chain all these preprocessing steps together in one line. For example, before feeding data into a Pytorch-based ResNet Executor, the image needs to be normalized and the color axis should be at first, not at the last. You can do this via:
@@ -150,7 +170,9 @@ d.chunks.plot_image_sprites('simpsons-chunks.png')
150170
Hmm, doesn't change so much. This is because we scan the whole image using sliding windows with no overlap (i.e. stride). Let's do a bit oversampling:
151171

152172
```python
153-
d.convert_image_tensor_to_sliding_windows(window_shape=(64, 64), strides=(10, 10), as_chunks=True)
173+
d.convert_image_tensor_to_sliding_windows(
174+
window_shape=(64, 64), strides=(10, 10), as_chunks=True
175+
)
154176
d.chunks.plot_image_sprites('simpsons-chunks-stride-10.png')
155177
```
156178

tests/unit/document/test_converters.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,18 @@ def test_uri_to_tensor():
5656
assert doc.mime_type == 'image/png'
5757

5858

59+
def test_uri_to_tensors_with_multi_page_tiff():
60+
doc = Document(uri=os.path.join(cur_dir, 'toydata/multi-page.tif'))
61+
doc.load_uri_to_image_tensor()
62+
63+
assert doc.tensor is None
64+
assert len(doc.chunks) == 3
65+
for chunk in doc.chunks:
66+
assert isinstance(chunk.tensor, np.ndarray)
67+
assert chunk.tensor.ndim == 3
68+
assert chunk.tensor.shape[-1] == 3
69+
70+
5971
def test_datauri_to_tensor():
6072
doc = Document(uri=os.path.join(cur_dir, 'toydata/test.png'))
6173
doc.convert_uri_to_datauri()
31 KB
Binary file not shown.

0 commit comments

Comments
 (0)