ADD: Transparency options for geotiff writer (#1496)

* ADD: Transparency options for geotiff writer

* FIX: Format input block

* FIX: Format import block

* FIX: PEP8 format fixes

* STY: Update PEP8 formatting.

---------

Co-authored-by: Zach Sherman <[email protected]>
Co-authored-by: zssherman <[email protected]>

pyart/io/output_to_geotiff.py

@@ -33,6 +33,8 @@ def write_grid_geotiff(
     warp=False,
     sld=False,
     use_doublequotes=True,
+    transparent_bg=True,
+    opacity=1.0,
 ):
     """
     Write a Py-ART Grid object to a GeoTIFF file.
@@ -100,6 +102,17 @@ def write_grid_geotiff(
 
         False - Use single quotes instead.
 
+    transparent_bg : bool, optional
+        True - Sets alpha value of masked pixels to zero producing a
+               transparent background
+
+        False - Sets alpha value of masked pixels to value assigned by
+                opacity parameter
+
+    opacity : float, optional
+        Alpha value to be assigned to all pixels (except for transparent background)
+        Value must be between 0 (transparent) and 1 (opaque)
+
     """
     if not IMPORT_FLAG:
         raise MissingOptionalDependency("GDAL not detected, GeoTIFF output failure!")
@@ -156,9 +169,11 @@ def write_grid_geotiff(
         )
     else:
         # Assign data RGB levels based on value relative to vmax/vmin
-        rarr, garr, barr = _get_rgb_values(data, vmin, vmax, color_levels, cmap)
+        rarr, garr, barr, aarr = _get_rgb_values(
+            data, vmin, vmax, color_levels, cmap, transparent_bg, opacity
+        )
         dst_ds = out_driver.Create(
-            ofile, data.shape[1], data.shape[0], 3, gdal.GDT_Byte
+            ofile, data.shape[1], data.shape[0], 4, gdal.GDT_Byte
         )
 
     # Common Projection and GeoTransform
@@ -172,6 +187,7 @@ def write_grid_geotiff(
         dst_ds.GetRasterBand(1).WriteArray(rarr[::-1, :])
         dst_ds.GetRasterBand(2).WriteArray(garr[::-1, :])
         dst_ds.GetRasterBand(3).WriteArray(barr[::-1, :])
+        dst_ds.GetRasterBand(4).WriteArray(aarr[::-1, :])
     dst_ds.FlushCache()
     dst_ds = None
 
@@ -202,7 +218,7 @@ def write_grid_geotiff(
         shutil.move(ofile + "_tmp.tif", ofile)
 
 
-def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
+def _get_rgb_values(data, vmin, vmax, color_levels, cmap, transpbg, op):
     """
     Get RGB values for later output to GeoTIFF, given a 2D data field,
     display min/max and color table info. Missing data get numpy.nan.
@@ -221,6 +237,11 @@ def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
         with steps << 255 (e.g., hydrometeor ID).
     cmap : str or matplotlib.colors.Colormap object, optional
         Colormap to use for RGB output or SLD file.
+    transpbg : bool
+        True - generate alpha channel with masked values set to 0
+        False - generate alpha channel with masked values set to op value
+    op : float
+        Opacity of image, value of 0 is transparent, value of 1 is opaque
 
     Returns
     -------
@@ -230,6 +251,8 @@ def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
         Blue channel indices (range = 0-255).
     garr : numpy.ndarray object, dtype int
         Green channel indices (range = 0-255).
+    aarr : numpy.ndarray object, dtype int
+        Alpha channel indices (range = 0-255).
 
     """
     frac = (data - vmin) / float(vmax - vmin)
@@ -242,6 +265,7 @@ def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
     rarr = []
     garr = []
     barr = []
+    aarr = []
     cmap = plt.get_cmap(cmap)
     for val in index:
         if not np.isnan(val):
@@ -250,14 +274,21 @@ def _get_rgb_values(data, vmin, vmax, color_levels, cmap):
             rarr.append(int(np.round(r * 255)))
             garr.append(int(np.round(g * 255)))
             barr.append(int(np.round(b * 255)))
+            aarr.append(int(np.round(op * 255)))
         else:
             rarr.append(np.nan)
             garr.append(np.nan)
             barr.append(np.nan)
+            if not transpbg:
+                aarr.append(int(np.round(op * 255)))
+            else:
+                aarr.append(0)
+
     rarr = np.reshape(rarr, data.shape)
     garr = np.reshape(garr, data.shape)
     barr = np.reshape(barr, data.shape)
-    return rarr, garr, barr
+    aarr = np.reshape(aarr, data.shape)
+    return rarr, garr, barr, aarr
 
 
 def _create_sld(cmap, vmin, vmax, filename, color_levels=None):

tests/io/test_output_to_geotiff.py

@@ -1,9 +1,11 @@
 """ Unit Tests for Py-ART's output_to_geotiff.py module. """
 
 import warnings
+from pathlib import Path
 
 import numpy as np
 import pytest
+from PIL import Image
 
 import pyart
 
@@ -16,8 +18,8 @@ def test__get_rgb_values_nan():
     data[5] = np.nan
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", category=RuntimeWarning)
-        rarr, barr, garr = pyart.io.output_to_geotiff._get_rgb_values(
-            data, 0, 10, None, "jet"
+        rarr, barr, garr, aarr = pyart.io.output_to_geotiff._get_rgb_values(
+            data, 0, 10, None, "jet", False, 1
         )
     assert np.isnan(rarr[5])
     assert np.isnan(barr[5])
@@ -53,6 +55,22 @@ def make_tiny_grid():
     return grid
 
 
+def make_tiny_grid_with_mask():
+    """Make a tiny grid."""
+    grid_shape = (2, 10, 8)
+    grid_limits = ((0, 500), (-400000, 400000), (-300000, 300000))
+    grid = pyart.testing.make_empty_grid(grid_shape, grid_limits)
+    fdata = np.zeros((2, 10, 8), dtype="float32")
+    fdata[:, 2:-2, 1:-1] = 10.0
+    fdata[:, 3:-3, 2:-2] = 20.0
+    fdata[:, 4:-4, 3:-3] = 30.0
+    fdata[1] += 5
+    fdata[fdata == 0] = np.nan
+    rdic = {"data": fdata, "long_name": "reflectivity", "units": "dBz"}
+    grid.fields = {"reflectivity": rdic}
+    return grid
+
+
 @pytest.mark.skipif(
     not pyart.io.output_to_geotiff.IMPORT_FLAG, reason="GDAL is not installed."
 )
@@ -132,3 +150,58 @@ def test_write_grid_geotiff_sld():
 def test_write_grid_geotiff_missing_field():
     grid = make_tiny_grid()
     pytest.raises(KeyError, pyart.io.write_grid_geotiff, grid, "test.foo", "foobar")
+
+
+@pytest.mark.skipif(
+    not pyart.io.output_to_geotiff.IMPORT_FLAG, reason="GDAL is not installed."
+)
+def test_write_grid_geotiff_transparent_background():
+    grid = make_tiny_grid_with_mask()
+
+    try:
+        with pyart.testing.InTemporaryDirectory() as tmpdir:
+            tmp = Path(tmpdir)
+            outname = tmp / "transparent_bg.tif"
+            pyart.io.write_grid_geotiff(
+                grid,
+                str(outname),
+                "reflectivity",
+                rgb=True,
+                cmap="pyart_HomeyerRainbow",
+                vmin=0,
+                vmax=40,
+                transparent_bg=True,
+                opacity=1,
+            )
+            imgname = outname.rename(tmp / "transparent_bg.tiff")
+            img = Image.open(imgname)
+            img.show()
+    except PermissionError:
+        pass
+
+
+@pytest.mark.skipif(
+    not pyart.io.output_to_geotiff.IMPORT_FLAG, reason="GDAL is not installed."
+)
+def test_write_grid_geotiff_opacity():
+    grid = make_tiny_grid_with_mask()
+    try:
+        with pyart.testing.InTemporaryDirectory() as tmpdir:
+            tmp = Path(tmpdir)
+            outname = tmp / "opacity.tif"
+            pyart.io.write_grid_geotiff(
+                grid,
+                str(outname),
+                "reflectivity",
+                rgb=True,
+                cmap="pyart_HomeyerRainbow",
+                vmin=0,
+                vmax=40,
+                transparent_bg=False,
+                opacity=0.25,
+            )
+            imgname = outname.rename(tmp / "opacity.tiff")
+            img = Image.open(imgname)
+            img.show()
+    except PermissionError:
+        pass