preprocess sample image

mrocklin · Aug 14, 2015 · 1f08bc2 · 1f08bc2
1 parent f3c5340
commit 1f08bc2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/dask/array/image.py b/dask/array/image.py
@@ -51,6 +51,8 @@ def imread(filename, imread=None, preprocess=None):
     name = 'imread-%s' % tokenize(filenames, map(os.path.getmtime, filenames))
 
     sample = imread(filenames[0])
+    if preprocess:
+        sample = preprocess(sample)
 
     keys = [(name, i) + (0,) * len(sample.shape) for i in range(len(filenames))]
     if preprocess:

diff --git a/dask/array/tests/test_image.py b/dask/array/tests/test_image.py
@@ -46,7 +46,7 @@ def imread2(fn):
 def test_preprocess():
     def preprocess(x):
         x[:] = 1
-        return x
+        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, 4), dtype='i1')).all()
+        assert (im.compute() == np.ones((4, 2, 3), dtype='i1')).all()