From 9d72cfa75ced54acbef8a03250ad6210dbeaae49 Mon Sep 17 00:00:00 2001 From: Qiusheng Wu Date: Sun, 11 Aug 2024 23:56:44 -0400 Subject: [PATCH] Improve array_to_image function (#869) --- leafmap/common.py | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/leafmap/common.py b/leafmap/common.py index 048482bda1..4cc4345891 100644 --- a/leafmap/common.py +++ b/leafmap/common.py @@ -10623,6 +10623,7 @@ def array_to_image( import numpy as np import rasterio import xarray as xr + import rioxarray from rasterio.transform import Affine if output is None: @@ -10640,14 +10641,31 @@ def array_to_image( ) if isinstance(array, xr.DataArray): - coords = [coord for coord in array.coords] - if coords[0] == "time": - x_dim = coords[1] - y_dim = coords[2] - array = ( - array.isel(time=0).rename({y_dim: "y", x_dim: "x"}).transpose("y", "x") + if ( + hasattr(array, "rio") + and (array.rio.crs is not None) + and (array.rio.transform() is not None) + ): + + if "latitude" in array.dims and "longitude" in array.dims: + array = array.rename({"latitude": "y", "longitude": "x"}) + elif "lat" in array.dims and "lon" in array.dims: + array = array.rename({"lat": "y", "lon": "x"}) + + if array.ndim == 2 and ("x" in array.dims) and ("y" in array.dims): + array = array.transpose("y", "x") + elif array.ndim == 3 and ("x" in array.dims) and ("y" in array.dims): + dims = list(array.dims) + dims.remove("x") + dims.remove("y") + array = array.transpose(dims[0], "y", "x") + if "long_name" in array.attrs: + array.attrs.pop("long_name") + + array.rio.to_raster( + output, driver=driver, compress=compress, dtype=dtype, **kwargs ) - array = array.values + return if array.ndim == 3 and transpose: array = np.transpose(array, (1, 2, 0))