-
-
Notifications
You must be signed in to change notification settings - Fork 1.7k
/
test_image.py
61 lines (43 loc) · 1.7 KB
/
test_image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from __future__ import annotations
import os
from contextlib import contextmanager
import pytest
from dask.array.numpy_compat import _numpy_200
if _numpy_200:
# `import skimage` fails with `numpy=2`
# xref https://github.com/scikit-image/scikit-image/pull/7118
pytest.skip("skimage doesn't yet support numpy=2", allow_module_level=True)
pytest.importorskip("skimage")
import numpy as np
from skimage.io import imsave
from dask.array.image import imread as da_imread
from dask.utils import tmpdir
@contextmanager
def random_images(n, shape):
with tmpdir() as dirname:
for i in range(n):
fn = os.path.join(dirname, "image.%d.png" % i)
x = np.random.randint(0, 255, size=shape).astype("u1")
imsave(fn, x, check_contrast=False)
yield os.path.join(dirname, "*.png")
def test_imread():
with random_images(4, (5, 6, 3)) as globstring:
im = da_imread(globstring)
assert im.shape == (4, 5, 6, 3)
assert im.chunks == ((1, 1, 1, 1), (5,), (6,), (3,))
assert im.dtype == "uint8"
assert im.compute().shape == (4, 5, 6, 3)
assert im.compute().dtype == "uint8"
def test_imread_with_custom_function():
def imread2(fn):
return np.ones((2, 3, 4), dtype="i1")
with random_images(4, (5, 6, 3)) as globstring:
im = da_imread(globstring, imread=imread2)
assert (im.compute() == np.ones((4, 2, 3, 4), dtype="u1")).all()
def test_preprocess():
def preprocess(x):
x[:] = 1
return x[:, :, 0]
with random_images(4, (2, 3, 4)) as globstring:
im = da_imread(globstring, preprocess=preprocess)
assert (im.compute() == np.ones((4, 2, 3), dtype="u1")).all()