true_color: added output name and attrs (#484)

* true_color: name of output data array

* output attrs

* docs

* added contrast as input param

* example

* Update xrspatial/multispectral.py

Co-authored-by: Giancarlo Castro <giancastrok@gmail.com>

* Update xrspatial/multispectral.py

Co-authored-by: Giancarlo Castro <giancastrok@gmail.com>

* fix true_color function docstring

Co-authored-by: Giancarlo Castro <giancastrok@gmail.com>

Author: thuydotm

xrspatial/multispectral.py

@@ -1528,14 +1528,14 @@ def ebbi(red_agg: xr.DataArray,
 
 
 @ngjit
-def _normalize_data_cpu(data, min_val, max_val, pixel_max):
+def _normalize_data_cpu(data, min_val, max_val, pixel_max, contrast):
     out = np.zeros_like(data)
     out[:] = np.nan
 
     range_val = max_val - min_val
     rows, cols = data.shape
 
-    c = 10
+    c = contrast
     th = .125
 
     # check range_val to avoid dividing by zero
@@ -1550,35 +1550,37 @@ def _normalize_data_cpu(data, min_val, max_val, pixel_max):
     return out
 
 
-def _normalize_data_numpy(data, pixel_max):
+def _normalize_data_numpy(data, pixel_max, contrast):
     min_val = np.nanmin(data)
     max_val = np.nanmax(data)
-    out = _normalize_data_cpu(data, min_val, max_val, pixel_max)
+    out = _normalize_data_cpu(data, min_val, max_val, pixel_max, contrast)
     return out
 
 
-def _normalize_data_dask(data, pixel_max):
+def _normalize_data_dask(data, pixel_max, contrast):
     min_val = da.nanmin(data)
     max_val = da.nanmax(data)
-    out = da.map_blocks(_normalize_data_cpu, data, min_val, max_val, pixel_max,
-                        meta=np.array(()))
+    out = da.map_blocks(
+        _normalize_data_cpu, data, min_val, max_val, pixel_max, contrast,
+        meta=np.array(())
+    )
     return out
 
 
-def _normalize_data_cupy(data, pixel_max):
+def _normalize_data_cupy(data, pixel_max, contrast):
     raise NotImplementedError('Not Supported')
 
 
-def _normalize_data_dask_cupy(data, pixel_max):
+def _normalize_data_dask_cupy(data, pixel_max, contrast):
     raise NotImplementedError('Not Supported')
 
 
-def _normalize_data(agg, pixel_max=255.0):
+def _normalize_data(agg, pixel_max=255.0, contrast=10):
     mapper = ArrayTypeFunctionMapping(numpy_func=_normalize_data_numpy,
                                       dask_func=_normalize_data_dask,
                                       cupy_func=_normalize_data_cupy,
                                       dask_cupy_func=_normalize_data_dask_cupy)
-    out = mapper(agg)(agg.data, pixel_max)
+    out = mapper(agg)(agg.data, pixel_max, contrast)
     return out
 
 
@@ -1609,56 +1611,107 @@ def _alpha(red, nodata=1):
     return out
 
 
-def _true_color_numpy(r, g, b, nodata):
+def _true_color_numpy(r, g, b, nodata, contrast):
     a = np.where(np.logical_or(np.isnan(r), r <= nodata), 0, 255)
 
     h, w = r.shape
     out = np.zeros((h, w, 4), dtype=np.uint8)
 
     pixel_max = 255
-    out[:, :, 0] = (_normalize_data(r, pixel_max)).astype(np.uint8)
-    out[:, :, 1] = (_normalize_data(g, pixel_max)).astype(np.uint8)
-    out[:, :, 2] = (_normalize_data(b, pixel_max)).astype(np.uint8)
+    out[:, :, 0] = (_normalize_data(r, pixel_max, contrast)).astype(np.uint8)
+    out[:, :, 1] = (_normalize_data(g, pixel_max, contrast)).astype(np.uint8)
+    out[:, :, 2] = (_normalize_data(b, pixel_max, contrast)).astype(np.uint8)
     out[:, :, 3] = a.astype(np.uint8)
     return out
 
 
-def _true_color_dask(r, g, b, nodata):
+def _true_color_dask(r, g, b, nodata, contrast):
     pixel_max = 255
-    red = (_normalize_data(r, pixel_max)).astype(np.uint8)
-    green = (_normalize_data(g, pixel_max)).astype(np.uint8)
-    blue = (_normalize_data(b, pixel_max)).astype(np.uint8)
+    red = (_normalize_data(r, pixel_max, contrast)).astype(np.uint8)
+    green = (_normalize_data(g, pixel_max, contrast)).astype(np.uint8)
+    blue = (_normalize_data(b, pixel_max, contrast)).astype(np.uint8)
 
     alpha = _alpha(r, nodata).astype(np.uint8)
 
     out = da.stack([red, green, blue, alpha], axis=-1)
     return out
 
 
-def _true_color_cupy(r, g, b, nodata):
+def _true_color_cupy(r, g, b, nodata, contrast):
     raise NotImplementedError('Not Supported')
 
 
-def _true_color_dask_cupy(r, g, b, nodata):
+def _true_color_dask_cupy(r, g, b, nodata, contrast):
     raise NotImplementedError('Not Supported')
 
 
-def true_color(r, g, b, nodata=1, name='true_color'):
+def true_color(r, g, b, nodata=1, contrast=10.0, name='true_color'):
+    """
+    Create true color composite from a combination of red, green and
+    blue bands satellite images.
+
+    Parameters
+    ----------
+    r : xarray.DataArray
+        2D array of red band data.
+    g : xarray.DataArray
+        2D array of green band data.
+    b : xarray.DataArray
+        2D array of blue band data.
+    nodata : int, float numeric value
+        Nodata value of input DataArrays.
+    contrast : int, float numeric value
+        Contrast controlling parameter for output image.
+    name : str, default='true_color'
+        Name of output DataArray.
+
+    Returns
+    -------
+    true_color_agg : xarray.DataArray of the same type as inputs
+        3D array true color image with dims of [y, x, band].
+        All output attributes are copied from red band image.
+
+
+    Examples
+    --------
+    .. plot::
+       :include-source:
+
+        import matplotlib.pyplot as plt
+        from xrspatial.multispectral import true_color
+        from xrspatial.datasets import get_data
+
+        # Open Example Data
+        data = get_data('sentinel-2')
+
+        red = data['Red']
+        green = data['Green']
+        blue = data['Blue']
+
+        # Generate ARVI Aggregate Array
+        true_color_img = true_color(r=red, g=green, b=blue)
+        true_color_img.plot.imshow()
+    """
+
     mapper = ArrayTypeFunctionMapping(numpy_func=_true_color_numpy,
                                       dask_func=_true_color_dask,
                                       cupy_func=_true_color_cupy,
                                       dask_cupy_func=_true_color_dask_cupy)
     with warnings.catch_warnings():
         warnings.simplefilter('ignore')
-        out = mapper(r)(r, g, b, nodata)
+        out = mapper(r)(r, g, b, nodata, contrast)
 
-    # TODO: output metadata: coords, dims, atts
+    # TODO: output metadata: coords, dims, attrs
     _dims = ['y', 'x', 'band']
+    _attrs = r.attrs
     _coords = {'y': r['y'],
                'x': r['x'],
                'band': [0, 1, 2, 3]}
 
-    return DataArray(out,
-                     dims=_dims,
-                     coords=_coords
-                     )
+    return DataArray(
+        out,
+        name=name,
+        dims=_dims,
+        coords=_coords,
+        attrs=_attrs,
+    )

xrspatial/tests/test_multispectral.py

@@ -788,13 +788,17 @@ def test_true_color_cpu():
     red = create_test_arr(arr1)
     green = create_test_arr(arr2)
     blue = create_test_arr(arr3)
-    numpy_result = true_color(red, green, blue)
+    numpy_result = true_color(red, green, blue, name='np_true_color')
+    assert numpy_result.name == 'np_true_color'
 
     # dask
     red_dask = create_test_arr(arr1, backend='dask')
     green_dask = create_test_arr(arr2, backend='dask')
     blue_dask = create_test_arr(arr3, backend='dask')
-    dask_result = true_color(red_dask, green_dask, blue_dask)
+    dask_result = true_color(
+        red_dask, green_dask, blue_dask, name='dask_true_color'
+    )
+    assert dask_result.name == 'dask_true_color'
 
     # TODO: test output metadata: dims, coords, attrs
     assert isinstance(numpy_result, xa.DataArray)