From ee84448a80c80cb15a1cb8607ecc37815f1a6551 Mon Sep 17 00:00:00 2001 From: Qiusheng Wu Date: Sat, 14 Dec 2024 14:52:37 -0500 Subject: [PATCH] Add support for raster custom colormap --- docs/notebooks/103_raster_colormap.ipynb | 206 +++++++++++++++++++++++ docs/tutorials.md | 2 + examples/README.md | 2 + leafmap/common.py | 144 +++++++++++++++- mkdocs.yml | 1 + 5 files changed, 348 insertions(+), 7 deletions(-) create mode 100644 docs/notebooks/103_raster_colormap.ipynb diff --git a/docs/notebooks/103_raster_colormap.ipynb b/docs/notebooks/103_raster_colormap.ipynb new file mode 100644 index 0000000000..28136f5213 --- /dev/null +++ b/docs/notebooks/103_raster_colormap.ipynb @@ -0,0 +1,206 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![image](https://jupyterlite.rtfd.io/en/latest/_static/badge.svg)](https://demo.leafmap.org/lab/index.html?path=notebooks/103_raster_colormap.ipynb)\n", + "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/leafmap/blob/master/docs/notebooks/103_raster_colormap.ipynb)\n", + "[![image](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/opengeos/leafmap/HEAD)\n", + "\n", + "**Applying a custom colormap to a raster dataset**\n", + "\n", + "Uncomment the following line to install the `leafmap` package." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "# %pip install -U \"leafmap[raster]\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import leafmap\n", + "import rioxarray as rxr" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Download a sample dataset from GitHub. This dataset is a GeoTIFF file containing the surface water extent in Las Vegas. This dataset is a [NASA OPERA DSWx](https://www.jpl.nasa.gov/go/opera/products/dswx-product-suite/) product. " + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "url = \"https://github.com/opengeos/datasets/releases/download/raster/OPERA_L3_DSWx_WTR.tif\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "filepath = leafmap.download_file(url, quiet=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the dataset as an xarray DataArray." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "da = rxr.open_rasterio(filepath)\n", + "# da" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The original raster file contains a colormap. We can get the colormap from the raster file as follows:" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "colormap = leafmap.get_image_colormap(filepath)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Alternatively, we can define a custom colormap as follows:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "colormap = {\n", + " 0: (255, 255, 255),\n", + " 1: (0, 0, 255),\n", + " 2: (180, 213, 244),\n", + " 252: (0, 255, 255),\n", + " 253: (175, 175, 175),\n", + " 254: (0, 0, 127),\n", + " 255: (0, 0, 0),\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can apply any data processing types to the xarray DataArray. After that, convert the xarray DataArray to an image in the memory and apply the custom colormap." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "image = leafmap.array_to_image(da, colormap=colormap)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Define a legend dictionary to display the legend." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "legend_dict = {\n", + " \"0: Not water\": (255, 255, 255),\n", + " \"1: Open water\": (0, 0, 255),\n", + " \"2: Partial surface water\": (180, 213, 244),\n", + " \"252: Snow/ice\": (0, 255, 255),\n", + " \"253: Cloud/cloud shadow\": (175, 175, 175),\n", + " \"254: Ocean masked\": (0, 0, 127),\n", + " \"255: Fill value (no data)\": (0, 0, 0),\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Visualize the raster dataset with the custom colormap and the legend." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = leafmap.Map()\n", + "m.add_basemap(\"HYBRID\")\n", + "m.add_raster(image, layer_name=\"Water\", nodata=255)\n", + "m.add_legend(legend_dict=legend_dict)\n", + "m" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "![image](https://github.com/user-attachments/assets/495c6e91-a722-4618-a2f7-3bbca64adca9)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "geo", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.2" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/tutorials.md b/docs/tutorials.md index 5cc9d692fe..6853740e3b 100644 --- a/docs/tutorials.md +++ b/docs/tutorials.md @@ -113,6 +113,8 @@ 99. Retrieving wetland boundaries from the National Wetlands Inventory (NWI) ([notebook](https://leafmap.org/notebooks/99_wetlands)) 100. Visualizing the National Land Cover Database (NLCD) data products with Leafmap ([notebook](https://leafmap.org/notebooks/100_nlcd)) 101. Searching and Visualizing NASA OPERA Data Products Interactively ([notebook](https://leafmap.org/notebooks/101_nasa_opera)) +102. Mapping Overture Buildings and Foursquare Places with Leafmap + Fused ([notebook](https://leafmap.org/notebooks/102_fused)) +103. Applying a custom colormap to a raster dataset ([notebook](https://leafmap.org/notebooks/103_raster_colormap)) ## Demo diff --git a/examples/README.md b/examples/README.md index 163c383637..2584657f43 100644 --- a/examples/README.md +++ b/examples/README.md @@ -120,6 +120,8 @@ 99. Retrieving wetland boundaries from the National Wetlands Inventory (NWI) ([notebook](https://leafmap.org/notebooks/99_wetlands)) 100. Visualizing the National Land Cover Database (NLCD) data products with Leafmap ([notebook](https://leafmap.org/notebooks/100_nlcd)) 101. Searching and Visualizing NASA OPERA Data Products Interactively ([notebook](https://leafmap.org/notebooks/101_nasa_opera)) +102. Mapping Overture Buildings and Foursquare Places with Leafmap + Fused ([notebook](https://leafmap.org/notebooks/102_fused)) +103. Applying a custom colormap to a raster dataset ([notebook](https://leafmap.org/notebooks/103_raster_colormap)) ## Demo diff --git a/leafmap/common.py b/leafmap/common.py index 05e38b7fb9..932582f6ce 100644 --- a/leafmap/common.py +++ b/leafmap/common.py @@ -10475,6 +10475,7 @@ def array_to_memory_file( crs: str = None, transform: tuple = None, driver="COG", + colormap: dict = None, **kwargs, ): """Convert a NumPy array to a memory file. @@ -10488,8 +10489,9 @@ def array_to_memory_file( cellsize (float, optional): The cell size of the array if source is not provided. Defaults to None. crs (str, optional): The coordinate reference system of the array if source is not provided. Defaults to None. transform (tuple, optional): The affine transformation matrix if source is not provided. - Can be rio.transform() or a tuple like (0.5, 0.0, -180.25, 0.0, -0.5, 83.780361). Defaults to None + Can be rio.transform() or a tuple like (0.5, 0.0, -180.25, 0.0, -0.5, 83.780361). Defaults to None. driver (str, optional): The driver to use for creating the output file, such as 'GTiff'. Defaults to "COG". + colormap (dict, optional): A dictionary defining the colormap (value: (R, G, B, A)). **kwargs: Additional keyword arguments to be passed to the rasterio.open() function. Returns: @@ -10510,8 +10512,7 @@ def array_to_memory_file( ) if hasattr(array, "rio"): if hasattr(array.rio, "crs"): - if array.rio.crs is not None: - crs = array.rio.crs + crs = array.rio.crs if transform is None and hasattr(array.rio, "transform"): transform = array.rio.transform() elif source is None: @@ -10575,7 +10576,6 @@ def array_to_memory_file( # Convert the array to the best dtype array = array.astype(dtype) - # Define the GeoTIFF metadata metadata = { "driver": driver, @@ -10601,12 +10601,15 @@ def array_to_memory_file( if array.ndim == 2: dst.write(array, 1) + if colormap: + dst.write_colormap(1, colormap) elif array.ndim == 3: for i in range(array.shape[2]): dst.write(array[:, :, i], i + 1) + if colormap: + dst.write_colormap(i + 1, colormap) dst.close() - # Read the dataset from memory dataset_reader = rasterio.open(dst.name, mode="r") @@ -10624,6 +10627,7 @@ def array_to_image( crs: str = None, transform: tuple = None, driver: str = "COG", + colormap: dict = None, **kwargs, ) -> str: """Save a NumPy array as a GeoTIFF using the projection information from an existing GeoTIFF file. @@ -10640,6 +10644,7 @@ def array_to_image( transform (tuple, optional): The affine transformation matrix, can be rio.transform() or a tuple like (0.5, 0.0, -180.25, 0.0, -0.5, 83.780361). Defaults to None. driver (str, optional): The driver to use for creating the output file, such as 'GTiff'. Defaults to "COG". + colormap (dict, optional): A dictionary defining the colormap (value: (R, G, B, A)). **kwargs: Additional keyword arguments to be passed to the rasterio.open() function. """ @@ -10660,6 +10665,7 @@ def array_to_image( crs=crs, transform=transform, driver=driver, + colormap=colormap, **kwargs, ) @@ -10688,7 +10694,9 @@ def array_to_image( array.rio.to_raster( output, driver=driver, compress=compress, dtype=dtype, **kwargs ) - return + if colormap: + write_image_colormap(output, colormap, output) + return output if array.ndim == 3 and transpose: array = np.transpose(array, (1, 2, 0)) @@ -10771,14 +10779,18 @@ def array_to_image( metadata["compress"] = compress metadata.update(**kwargs) - # Create a new GeoTIFF file and write the array to it with rasterio.open(output, "w", **metadata) as dst: if array.ndim == 2: dst.write(array, 1) + if colormap: + dst.write_colormap(1, colormap) elif array.ndim == 3: for i in range(array.shape[2]): dst.write(array[:, :, i], i + 1) + if colormap: + dst.write_colormap(i + 1, colormap) + return output def images_to_tiles( @@ -15563,3 +15575,121 @@ def download_mapillary_images( download_mapillary_image( image_id=image_id, output=output, resolution=resolution, **kwargs ) + + +def get_image_colormap(image, index=1): + """ + Retrieve the colormap from an image. + + Args: + image (str, rasterio.io.DatasetReader, rioxarray.DataArray): + The input image. It can be: + - A file path to a raster image (string). + - A rasterio dataset. + - A rioxarray DataArray. + index (int): The band index to retrieve the colormap from (default is 1). + + Returns: + dict: A dictionary representing the colormap (value: (R, G, B, A)), or None if no colormap is found. + + Raises: + ValueError: If the input image type is unsupported. + """ + import rasterio + import rioxarray + import xarray as xr + + dataset = None + + if isinstance(image, str): # File path + with rasterio.open(image) as ds: + return ds.colormap(index) if ds.count > 0 else None + elif isinstance(image, rasterio.io.DatasetReader): # rasterio dataset + dataset = image + elif isinstance(image, xr.DataArray) or isinstance(image, xr.Dataset): + source = image.encoding.get("source") + if source: + with rasterio.open(source) as ds: + return ds.colormap(index) if ds.count > 0 else None + else: + raise ValueError( + "Cannot extract colormap: DataArray does not have a source." + ) + else: + raise ValueError( + "Unsupported input type. Provide a file path, rasterio dataset, or rioxarray DataArray." + ) + + if dataset: + return dataset.colormap(index) if dataset.count > 0 else None + + +def write_image_colormap(image, colormap, output_path=None): + """ + Apply or update a colormap to a raster image. + + Args: + image (str, rasterio.io.DatasetReader, rioxarray.DataArray): + The input image. It can be: + - A file path to a raster image (string). + - A rasterio dataset. + - A rioxarray DataArray. + colormap (dict): A dictionary defining the colormap (value: (R, G, B, A)). + output_path (str, optional): Path to save the updated raster image. + If None, the original file is updated in-memory. + + Returns: + str: Path to the updated raster image. + + Raises: + ValueError: If the input image type is unsupported. + """ + import rasterio + import rioxarray + import xarray as xr + + dataset = None + src_profile = None + src_data = None + + if isinstance(image, str): # File path + with rasterio.open(image) as ds: + dataset = ds + src_profile = ds.profile + src_data = ds.read(1) # Assuming single-band + elif isinstance(image, rasterio.io.DatasetReader): # rasterio dataset + dataset = image + src_profile = dataset.profile + src_data = dataset.read(1) # Assuming single-band + elif isinstance(image, xr.DataArray): # rioxarray DataArray + source = image.encoding.get("source") + if source: + with rasterio.open(source) as ds: + dataset = ds + src_profile = ds.profile + src_data = ds.read(1) # Assuming single-band + else: + raise ValueError("Cannot apply colormap: DataArray does not have a source.") + else: + raise ValueError( + "Unsupported input type. Provide a file path, rasterio dataset, or rioxarray DataArray." + ) + + # Ensure the dataset is single-band + if dataset.count != 1: + raise ValueError( + "Colormaps can only be applied to single-band raster datasets." + ) + + # Update the profile and colormap + src_profile.update(dtype=src_data.dtype, count=1) + + if not output_path: + output_path = "output_with_colormap.tif" + + # Write the updated dataset with the colormap + with rasterio.open(output_path, "w", **src_profile) as dst: + dst.write(src_data, 1) + dst.write_colormap(1, colormap) + + return output_path diff --git a/mkdocs.yml b/mkdocs.yml index 1abf46a59f..2c88531b3e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -340,3 +340,4 @@ nav: - notebooks/100_nlcd.ipynb - notebooks/101_nasa_opera.ipynb - notebooks/102_fused.ipynb + - notebooks/103_raster_colormap.ipynb