added cupy support to true_color function (#866)

Author: brendancol

README.md

@@ -176,7 +176,7 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
 | [Normalized Difference Vegetation Index (NDVI)](xrspatial/multispectral.py) | Quantifies vegetation density from red and NIR band difference | ✅️ |✅️ | ✅️ |✅️ |
 | [Soil Adjusted Vegetation Index (SAVI)](xrspatial/multispectral.py) | Vegetation index with soil brightness correction factor | ✅️ |✅️ | ✅️ |✅️ |
 | [Structure Insensitive Pigment Index (SIPI)](xrspatial/multispectral.py) | Estimates carotenoid-to-chlorophyll ratio for plant stress detection | ✅️ |✅️ | ✅️ |✅️ |
-| [True Color](xrspatial/multispectral.py) | Composites red, green, and blue bands into a natural color image | ✅️ | ️ | ✅️ | ️ |
+| [True Color](xrspatial/multispectral.py) | Composites red, green, and blue bands into a natural color image | ✅️ | ✅️ | ✅️ | ✅️ |
 
 -------
 

xrspatial/multispectral.py

@@ -1371,11 +1371,35 @@ def _normalize_data_dask(data, pixel_max, c, th):
 
 
 def _normalize_data_cupy(data, pixel_max, c, th):
-    raise NotImplementedError('Not Supported')
+    min_val = cupy.nanmin(data)
+    max_val = cupy.nanmax(data)
+    range_val = max_val - min_val
+    out = cupy.full(data.shape, cupy.nan, dtype=cupy.float32)
+    if range_val != 0:
+        norm = (data - min_val) / range_val
+        norm = 1 / (1 + cupy.exp(c * (th - norm)))
+        out = norm * pixel_max
+    return out
+
+
+def _normalize_data_cupy_block(data, min_val, max_val, pixel_max, c, th):
+    range_val = max_val - min_val
+    out = cupy.full(data.shape, cupy.nan, dtype=cupy.float32)
+    if range_val != 0:
+        norm = (data - min_val) / range_val
+        norm = 1 / (1 + cupy.exp(c * (th - norm)))
+        out = norm * pixel_max
+    return out
 
 
 def _normalize_data_dask_cupy(data, pixel_max, c, th):
-    raise NotImplementedError('Not Supported')
+    min_val = da.nanmin(data)
+    max_val = da.nanmax(data)
+    out = da.map_blocks(
+        _normalize_data_cupy_block, data, min_val, max_val, pixel_max,
+        c, th, meta=cupy.array(())
+    )
+    return out
 
 
 def _normalize_data(agg, pixel_max, c, th):
@@ -1416,6 +1440,32 @@ def _true_color_dask(r, g, b, nodata, c, th):
     return out
 
 
+def _true_color_cupy(r, g, b, nodata, c, th):
+    pixel_max = 255
+    r_data = r.data
+    a = cupy.where(
+        cupy.logical_or(cupy.isnan(r_data), r_data <= nodata), 0, pixel_max
+    ).astype(cupy.uint8)
+    red = (_normalize_data(r, pixel_max, c, th)).astype(cupy.uint8)
+    green = (_normalize_data(g, pixel_max, c, th)).astype(cupy.uint8)
+    blue = (_normalize_data(b, pixel_max, c, th)).astype(cupy.uint8)
+    out = cupy.stack([red, green, blue, a], axis=-1)
+    return out
+
+
+def _true_color_dask_cupy(r, g, b, nodata, c, th):
+    pixel_max = 255
+    r_data = r.data
+    alpha = da.where(
+        da.logical_or(da.isnan(r_data), r_data <= nodata), 0, pixel_max
+    ).astype(cupy.uint8)
+    red = (_normalize_data(r, pixel_max, c, th)).astype(cupy.uint8)
+    green = (_normalize_data(g, pixel_max, c, th)).astype(cupy.uint8)
+    blue = (_normalize_data(b, pixel_max, c, th)).astype(cupy.uint8)
+    out = da.stack([red, green, blue, alpha], axis=-1)
+    return out
+
+
 def true_color(r, g, b, nodata=1, c=10.0, th=0.125, name='true_color'):
     """
     Create true color composite from a combination of red, green and
@@ -1468,12 +1518,8 @@ def true_color(r, g, b, nodata=1, c=10.0, th=0.125, name='true_color'):
     mapper = ArrayTypeFunctionMapping(
         numpy_func=_true_color_numpy,
         dask_func=_true_color_dask,
-        cupy_func=lambda *args: not_implemented_func(
-            *args, messages='true_color() does not support cupy backed DataArray',  # noqa
-        ),
-        dask_cupy_func=lambda *args: not_implemented_func(
-            *args, messages='true_color() does not support dask with cupy backed DataArray',  # noqa
-        ),
+        cupy_func=_true_color_cupy,
+        dask_cupy_func=_true_color_dask_cupy,
     )
     with warnings.catch_warnings():
         warnings.simplefilter('ignore')

xrspatial/tests/test_multispectral.py

@@ -613,3 +613,34 @@ def test_true_color_numpy_equals_dask_numpy(random_data):
     np.testing.assert_allclose(
         numpy_result.data, dask_result.compute().data, equal_nan=True
     )
+
+
+@cuda_and_cupy_available
+@pytest.mark.parametrize("size", [(2, 4), (10, 15)])
+@pytest.mark.parametrize(
+    "dtype", [np.int32, np.int64, np.uint32, np.uint64, np.float32, np.float64])
+@pytest.mark.parametrize("backend", ["cupy", "dask+cupy"])
+def test_true_color_gpu(random_data, backend):
+    # numpy baseline
+    red_np = create_test_raster(random_data, backend="numpy")
+    green_np = create_test_raster(random_data, backend="numpy")
+    blue_np = create_test_raster(random_data, backend="numpy")
+    numpy_result = true_color(red_np, green_np, blue_np)
+
+    # gpu version
+    red_gpu = create_test_raster(random_data, backend=backend)
+    green_gpu = create_test_raster(random_data, backend=backend)
+    blue_gpu = create_test_raster(random_data, backend=backend)
+    gpu_result = true_color(red_gpu, green_gpu, blue_gpu)
+
+    general_output_checks(red_gpu, gpu_result, verify_attrs=False)
+
+    gpu_data = gpu_result.data
+    if hasattr(gpu_data, 'compute'):
+        gpu_data = gpu_data.compute()
+    if hasattr(gpu_data, 'get'):
+        gpu_data = gpu_data.get()
+
+    np.testing.assert_allclose(
+        numpy_result.data, gpu_data, equal_nan=True
+    )