From 14f625a10633e92c12818c41ea9269322b25e1ee Mon Sep 17 00:00:00 2001 From: qubitqualia <39470632+qubitqualia@users.noreply.github.com> Date: Tue, 24 Sep 2024 13:41:30 -0400 Subject: [PATCH] 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 <19153455+zssherman@users.noreply.github.com> Co-authored-by: zssherman --- pyart/io/output_to_geotiff.py | 39 +++++++++++++-- tests/io/test_output_to_geotiff.py | 77 +++++++++++++++++++++++++++++- 2 files changed, 110 insertions(+), 6 deletions(-) diff --git a/pyart/io/output_to_geotiff.py b/pyart/io/output_to_geotiff.py index 2f5262fa0e4..7e327dfa8e9 100644 --- a/pyart/io/output_to_geotiff.py +++ b/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): diff --git a/tests/io/test_output_to_geotiff.py b/tests/io/test_output_to_geotiff.py index 94c91c1a00e..7f9259b0e12 100644 --- a/tests/io/test_output_to_geotiff.py +++ b/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