Fix Image.clip_to_dtype and Image.astype with Dask arrays

giatools/image.py

@@ -417,7 +417,7 @@ def astype(
 
         # Special case: Conversion to `bool`
         elif dtype == bool:
-            labels = _np.unique(self.data)
+            labels = _unique(self.data)
             if len(labels) > 2:
                 raise ValueError(
                     f'Cannot convert image data from {self.data.dtype} to bool without overflows '
@@ -442,8 +442,7 @@ def astype(
                 dtype = resolve_unsignedinteger_to
 
             # Check for overflows
-            src_min = self.data.min().item()  # convert to native Python type (int, float)
-            src_max = self.data.max().item()  # convert to native Python type (int, float)
+            src_min, src_max = _get_min_max_values(self.data)
             if _np.issubdtype(dtype, _np.integer):
                 dst_min = _np.iinfo(dtype).min
                 dst_max = _np.iinfo(dtype).max
@@ -485,8 +484,7 @@ def clip_to_dtype(self, dtype: _np.dtype, force_copy: bool = False) -> _T.Self:
             raise TypeError('Clipping to boolean dtype is not supported.')
 
         # Determine the actual range of the source image
-        min_src_value = self.data.min().item()  # convert to native Python type (float, int)
-        max_src_value = self.data.max().item()  # convert to native Python type (float, int)
+        min_src_value, max_src_value = _get_min_max_values(self.data)
 
         # Determine the valid range for the target dtype
         if _np.issubdtype(dtype, _np.integer):
@@ -511,3 +509,26 @@ def clip_to_dtype(self, dtype: _np.dtype, force_copy: bool = False) -> _T.Self:
                 original_axes=self.original_axes,
                 metadata=self.metadata,
             )
+
+
+def _get_min_max_values(array: _T.NDArray) -> _T.Tuple[_T.Union[float, int], _T.Union[float, int]]:
+    if hasattr(array, 'compute'):  # Dask array
+        import dask.array as da
+        min_src_value, max_src_value = (
+            value.item()  # convert to native Python type (float, int)
+            for value in da.compute(array.min(), array.max())
+        )
+    else:  # NumPy array
+        min_src_value, max_src_value = (
+            value.item()  # convert to native Python type (float, int)
+            for value in (array.min(), array.max())
+        )
+    return min_src_value, max_src_value
+
+
+def _unique(array: _T.NDArray) -> _T.NDArray:
+    if hasattr(array, 'compute'):  # Dask array
+        import dask.array as da
+        return da.unique(array).compute()
+    else:  # NumPy array
+        return _np.unique(array)

tests/module/test__image.py

@@ -10,6 +10,7 @@
 
 import giatools.image
 from giatools.typing import (
+    Any,
     Optional,
     Tuple,
 )
@@ -19,6 +20,21 @@
     permute_axes,
 )
 
+exact_dtype_list = [
+    np.uint8,
+    np.int8,
+    np.uint16,
+    np.int16,
+    np.uint32,
+    np.int32,
+    np.uint64,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+    bool,
+]
+
 
 class ImageTestCase(unittest.TestCase):
 
@@ -269,22 +285,7 @@ def test__dask_array__zyx__iterate__yx(self, joint_axes: str):
 
 class ImageTestCase__dtype_mixin:
 
-    exact_dtype_list = [
-        np.uint8,
-        np.int8,
-        np.uint16,
-        np.int16,
-        np.uint32,
-        np.int32,
-        np.uint64,
-        np.int64,
-        np.float16,
-        np.float32,
-        np.float64,
-        bool,
-    ]
-
-    def create_non_bool_image(
+    def create_random_non_bool_image(
         self,
         dtype: np.dtype,
         shape=(2, 3, 26, 32),
@@ -320,10 +321,53 @@ def create_non_bool_image(
             metadata=metadata if metadata is not None else unittest.mock.Mock(),
         )
 
+    def create_random_bool_image(self) -> giatools.image.Image:
+        return giatools.image.Image(
+            data=np.random.choice(a=[False, True], size=(2, 3, 26, 32)),
+            axes='CYXZ',
+        )
 
-class Image__astype(ImageTestCase, ImageTestCase__dtype_mixin):
+    def create_const_value_image(
+        self,
+        dtype: np.dtype,
+        value: Any,
+        shape=(2, 3, 26, 32),
+        axes='CYXZ',
+        original_axes='QTCYXZ',
+        metadata: Optional[unittest.mock.Mock] = None,
+    ) -> giatools.image.Image:
+        return giatools.image.Image(
+            data=np.full(shape, value, dtype=dtype),
+            axes=axes,
+            original_axes=original_axes,
+            metadata=metadata if metadata is not None else unittest.mock.Mock(),
+        )
 
-    exact_dtype_list = list(frozenset(ImageTestCase__dtype_mixin.exact_dtype_list) - {bool})
+
+class ImageTestCase__dtype_mixin__dask(ImageTestCase__dtype_mixin):
+
+    def create_random_non_bool_image(self, *args, **kwargs) -> giatools.image.Image:
+        import dask.array as da
+        img = super().create_random_non_bool_image(*args, **kwargs)
+        img.data = da.from_array(img.data, chunks=(10,) * img.data.ndim)
+        return img
+
+    def create_random_bool_image(self) -> giatools.image.Image:
+        import dask.array as da
+        img = super().create_random_bool_image()
+        img.data = da.from_array(img.data, chunks=(10,) * img.data.ndim)
+        return img
+
+    def create_const_value_image(self, *args, **kwargs) -> giatools.image.Image:
+        import dask.array as da
+        img = super().create_const_value_image(*args, **kwargs)
+        img.data = da.from_array(img.data, chunks=(10,) * img.data.ndim)
+        return img
+
+
+class Image__astype__mixin:
+
+    exact_non_bool_dtype_list = list(frozenset(exact_dtype_list) - {bool})
 
     inexact_dtype_list = [
         np.floating,
@@ -346,7 +390,7 @@ def _test_non_bool_conversion(
             expected_dtype = dst_dtype
 
         # Create test image
-        img = self.create_non_bool_image(src_dtype)
+        img = self.create_random_non_bool_image(src_dtype)
         original_dtype = img.data.dtype
         original_metadata = img.metadata
         original_axes = img.axes
@@ -399,7 +443,7 @@ def _test_non_bool_conversion(
                     np.float16, np.float32, np.float64,
                 ) else np.iinfo(expected_dtype).max
             )
-            fallback_img = self.create_non_bool_image(
+            fallback_img = self.create_random_non_bool_image(
                 src_dtype,
                 min_value=0,
                 max_value=min((max_dst_value, max_src_value)),
@@ -437,8 +481,8 @@ def convert(_img):
                 self.assertIs(img.data, img_converted.data)
 
     def test__non_bool__exact(self):
-        for src_dtype in self.exact_dtype_list:
-            for dst_dtype in self.exact_dtype_list:
+        for src_dtype in self.exact_non_bool_dtype_list:
+            for dst_dtype in self.exact_non_bool_dtype_list:
                 for force_copy in (False, True):
                     with self.subTest(f'from {src_dtype} to {dst_dtype} (force_copy={force_copy})'):
                         self._test_non_bool_conversion(src_dtype, dst_dtype, force_copy=force_copy)
@@ -475,7 +519,7 @@ def _get_expected_dtype(src_dtype: np.dtype, dst_dtype: npt.DTypeLike) -> np.dty
             return dst_dtype
 
     def test__non_bool__inexact(self):
-        for src_dtype in self.exact_dtype_list:
+        for src_dtype in self.exact_non_bool_dtype_list:
             for dst_dtype in self.inexact_dtype_list:
                 for force_copy in (False, True):
                     with self.subTest(f'from {src_dtype} to {dst_dtype} (force_copy={force_copy})'):
@@ -494,7 +538,7 @@ def test__conversion__from_bool(self):
             original_axes='QTCYXZ',
         )
         assert img.data.dtype == bool  # sanity check
-        for dst_dtype in self.exact_dtype_list + self.inexact_dtype_list:
+        for dst_dtype in self.exact_non_bool_dtype_list + self.inexact_dtype_list:
             for force_copy in (False, True):
                 with self.subTest(f'from bool to {dst_dtype} (force_copy={force_copy})'):
                     expected_dtype = self._get_expected_dtype(bool, dst_dtype)
@@ -514,40 +558,74 @@ def _test_conversion_to_bool(self, img, expected_data):
                 np.testing.assert_array_equal(img_converted.data, expected_data)
 
     def test__conversion__to_bool__2_labels(self):
-        for src_dtype in self.exact_dtype_list:
-            img = self.create_non_bool_image(src_dtype)
+        for src_dtype in self.exact_non_bool_dtype_list:
+            img = self.create_random_non_bool_image(src_dtype)
             img.data = 10 + 5 * (img.data > img.data.mean()).astype(img.data.dtype)
-            assert list(np.unique(img.data)) == [10, 15]  # sanity check
+            assert list(np.unique(np.asarray(img.data))) == [10, 15]  # sanity check
             self._test_conversion_to_bool(img, expected_data=img.data > 12)
 
     def test__conversion__to_bool__1_non_zero_label(self):
-        for src_dtype in self.exact_dtype_list:
-            img = self.create_non_bool_image(src_dtype)
-            img.data.fill(15)
+        for src_dtype in self.exact_non_bool_dtype_list:
+            img = self.create_const_value_image(src_dtype, value=15)
             assert img.data.dtype == src_dtype  # sanity check
             self._test_conversion_to_bool(img, expected_data=np.ones(img.data.shape, dtype=bool))
 
     def test__conversion__to_bool__1_zero_label(self):
-        for src_dtype in self.exact_dtype_list:
-            img = self.create_non_bool_image(src_dtype)
-            img.data.fill(0)
+        for src_dtype in self.exact_non_bool_dtype_list:
+            img = self.create_const_value_image(src_dtype, value=0)
             assert img.data.dtype == src_dtype  # sanity check
             self._test_conversion_to_bool(img, expected_data=np.zeros(img.data.shape, dtype=bool))
 
     def test__conversion__to_bool__invalid(self):
-        for src_dtype in self.exact_dtype_list:
-            img = self.create_non_bool_image(src_dtype)
+        for src_dtype in self.exact_non_bool_dtype_list:
+            img = self.create_random_non_bool_image(src_dtype)
             img.data = np.random.randint(0, 3, img.data.shape).astype(src_dtype)
             assert len(np.unique(img.data)) > 2  # sanity check
             with self.subTest(f'from {src_dtype} to bool (invalid case)'):
                 with self.assertRaises(ValueError):
                     img.astype(bool)
 
 
-class Image__clip_to_dtype(ImageTestCase, ImageTestCase__dtype_mixin):
+class Image__astype(ImageTestCase, ImageTestCase__dtype_mixin, Image__astype__mixin):
+    pass  # Tests with NumPy arrays
+
+
+class Image__astype__dask(ImageTestCase, ImageTestCase__dtype_mixin__dask, Image__astype__mixin):
+    pass  # Tests with Dask arrays
+
+    @minimum_python_version(3, 11)
+    def test__non_bool__exact(self):
+        super().test__non_bool__exact()
+
+    @minimum_python_version(3, 11)
+    def test__non_bool__inexact(self):
+        super().test__non_bool__inexact()
+
+    @minimum_python_version(3, 11)
+    def test__conversion__from_bool(self):
+        super().test__conversion__from_bool()
+
+    @minimum_python_version(3, 11)
+    def test__conversion__to_bool__2_labels(self):
+        super().test__conversion__to_bool__2_labels()
+
+    @minimum_python_version(3, 11)
+    def test__conversion__to_bool__1_non_zero_label(self):
+        super().test__conversion__to_bool__1_non_zero_label()
+
+    @minimum_python_version(3, 11)
+    def test__conversion__to_bool__1_zero_label(self):
+        super().test__conversion__to_bool__1_zero_label()
+
+    @minimum_python_version(3, 11)
+    def test__conversion__to_bool__invalid(self):
+        super().test__conversion__to_bool__invalid()
+
+
+class Image__clip_to_dtype__mixin:
 
     def test__float32_to_int8__no_clip(self):
-        img = self.create_non_bool_image(
+        img = self.create_random_non_bool_image(
             dtype=np.float32,
             min_value=-20.0,
             max_value=+20.0,
@@ -557,7 +635,7 @@ def test__float32_to_int8__no_clip(self):
         self.assertIs(img_clipped, img)
 
     def test__float32_to_float16__clip_below(self):
-        img = self.create_non_bool_image(
+        img = self.create_random_non_bool_image(
             dtype=np.float32,
             min_value=-1e5,
             max_value=+1e2,
@@ -570,7 +648,7 @@ def test__float32_to_float16__clip_below(self):
         self.assertEqual(img_clipped.data.max(), +1e2)
 
     def test__float32_to_int8__clip_above(self):
-        img = self.create_non_bool_image(
+        img = self.create_random_non_bool_image(
             dtype=np.float32,
             min_value=-100,
             max_value=+1e5,
@@ -583,9 +661,30 @@ def test__float32_to_int8__clip_above(self):
         self.assertEqual(img_clipped.data.max(), +127.0)
 
     def test__bool_to_uint8(self):
-        img = giatools.image.Image(
-            data=np.random.choice(a=[False, True], size=(2, 3, 26, 32)),
-            axes='CYXZ',
-        )
+        img = self.create_random_bool_image()
         img_clipped = img.clip_to_dtype(np.uint8)
         self.assertIs(img_clipped, img)
+
+
+class Image__clip_to_dtype(ImageTestCase, ImageTestCase__dtype_mixin, Image__clip_to_dtype__mixin):
+    pass  # Tests with NumPy arrays
+
+
+class Image__clip_to_dtype__dask(ImageTestCase, ImageTestCase__dtype_mixin__dask, Image__clip_to_dtype__mixin):
+    pass  # Tests with Dask arrays
+
+    @minimum_python_version(3, 11)
+    def test__float32_to_int8__no_clip(self):
+        super().test__float32_to_int8__no_clip()
+
+    @minimum_python_version(3, 11)
+    def test__float32_to_float16__clip_below(self):
+        super().test__float32_to_float16__clip_below()
+
+    @minimum_python_version(3, 11)
+    def test__float32_to_int8__clip_above(self):
+        super().test__float32_to_int8__clip_above()
+
+    @minimum_python_version(3, 11)
+    def test__bool_to_uint8(self):
+        super().test__bool_to_uint8()

tests/unit/test__image.py

@@ -33,6 +33,12 @@ def setUp(self):
         self._np = unittest.mock.patch(
             'giatools.image._np'
         ).start()
+        self._get_min_max_values = unittest.mock.patch(
+            'giatools.image._get_min_max_values'
+        ).start()
+        self._unique = unittest.mock.patch(
+            'giatools.image._unique'
+        ).start()
 
         self.addCleanup(unittest.mock.patch.stopall)
 
@@ -194,8 +200,7 @@ def test__bool(self):
             self.img1.clip_to_dtype(bool)
 
     def test__to_superset_int(self):
-        self.img1.data.min.return_value.item.return_value = -15
-        self.img1.data.max.return_value.item.return_value = +15
+        self._get_min_max_values.return_value = (-15, +15)
         self._np.issubdtype.return_value = True  # target dtype is an integer type
         self._np.iinfo.return_value.min = -15
         self._np.iinfo.return_value.max = +15
@@ -204,8 +209,7 @@ def test__to_superset_int(self):
         self.img1.data.copy.assert_not_called()
 
     def test__to_superset_float(self):
-        self.img1.data.min.return_value.item.return_value = -15
-        self.img1.data.max.return_value.item.return_value = +15
+        self._get_min_max_values.return_value = (-15, +15)
         self._np.issubdtype.return_value = False  # target dtype is a float type
         self._np.finfo.return_value.min.item.return_value = -15.
         self._np.finfo.return_value.max.item.return_value = +15.