Fix array_to_image bug

opengeos · Jan 15, 2024 · 62aebd1 · 62aebd1
1 parent cb48ba8
commit 62aebd1
Show file tree

Hide file tree

Showing 3 changed files with 206 additions and 10 deletions.
diff --git a/docs/notebooks/89_image_array_viz.ipynb b/docs/notebooks/89_image_array_viz.ipynb
@@ -107,6 +107,92 @@
     "m"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also specify the image metadata (e.g., cellsize, crs, and transform) when creating the in-memory raster dataset.\n",
+    "\n",
+    "First, check the metadata of the origina image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.profile"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Check the crs of the original image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.crs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Check the transform of the original image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.transform"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create an in-memory raster dataset from the NDVI array and specify the cellsize, crs, and transform."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = (30.0, 0.0, -13651650.0, 0.0, -30.0, 4576290.0)\n",
+    "ndvi_image = leafmap.array_to_image(ndvi, cellsize=30, crs=\"EPSG:3857\", transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Add the NDVI image to the map."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = leafmap.Map()\n",
+    "m.add_raster(satellite, band=[1, 2, 3], nodata=-1, layer_name=\"Landsat 7\")\n",
+    "m.add_raster(ndvi_image, cmap=\"Greens\", layer_name=\"NDVI\")\n",
+    "m"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -193,7 +279,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

diff --git a/examples/notebooks/89_image_array_viz.ipynb b/examples/notebooks/89_image_array_viz.ipynb
@@ -107,6 +107,92 @@
     "m"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also specify the image metadata (e.g., cellsize, crs, and transform) when creating the in-memory raster dataset.\n",
+    "\n",
+    "First, check the metadata of the origina image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.profile"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Check the crs of the original image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.crs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Check the transform of the original image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset.transform"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create an in-memory raster dataset from the NDVI array and specify the cellsize, crs, and transform."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = (30.0, 0.0, -13651650.0, 0.0, -30.0, 4576290.0)\n",
+    "ndvi_image = leafmap.array_to_image(ndvi, cellsize=30, crs=\"EPSG:3857\", transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Add the NDVI image to the map."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = leafmap.Map()\n",
+    "m.add_raster(satellite, band=[1, 2, 3], nodata=-1, layer_name=\"Landsat 7\")\n",
+    "m.add_raster(ndvi_image, cmap=\"Greens\", layer_name=\"NDVI\")\n",
+    "m"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -193,7 +279,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

diff --git a/leafmap/common.py b/leafmap/common.py
@@ -10266,7 +10266,8 @@ def array_to_memory_file(
         transpose (bool, optional): Whether to transpose the array from (bands, rows, columns) to (rows, columns, bands). Defaults to True.
         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. 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
         driver (str, optional): The driver to use for creating the output file, such as 'GTiff'. Defaults to "COG".
         **kwargs: Additional keyword arguments to be passed to the rasterio.open() function.
 
@@ -10276,6 +10277,7 @@ def array_to_memory_file(
     import rasterio
     import numpy as np
     import xarray as xr
+    from rasterio.transform import Affine
 
     if isinstance(array, xr.DataArray):
         coords = [coord for coord in array.coords]
@@ -10304,7 +10306,7 @@ def array_to_memory_file(
                 "crs must be provided if source is not provided, such as EPSG:3857"
             )
 
-        if "transform" not in kwargs:
+        if transform is None:
             # Define the geotransformation parameters
             xmin, ymin, xmax, ymax = (
                 0,
@@ -10316,8 +10318,12 @@ def array_to_memory_file(
             transform = rasterio.transform.from_bounds(
                 xmin, ymin, xmax, ymax, array.shape[1], array.shape[0]
             )
-        else:
-            transform = kwargs["transform"]
+        elif isinstance(transform, Affine):
+            pass
+        elif isinstance(transform, (tuple, list)):
+            transform = Affine(*transform)
+
+        kwargs["transform"] = transform
 
     if dtype is None:
         # Determine the minimum and maximum values in the array
@@ -10386,6 +10392,7 @@ def array_to_image(
     transpose: bool = True,
     cellsize: float = None,
     crs: str = None,
+    transform: tuple = None,
     driver: str = "COG",
     **kwargs,
 ) -> str:
@@ -10400,17 +10407,29 @@ def array_to_image(
         transpose (bool, optional): Whether to transpose the array from (bands, rows, columns) to (rows, columns, bands). Defaults to True.
         cellsize (float, optional): The resolution of the output image in meters. Defaults to None.
         crs (str, optional): The CRS of the output image. Defaults to None.
+        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".
         **kwargs: Additional keyword arguments to be passed to the rasterio.open() function.
     """
 
     import numpy as np
     import rasterio
     import xarray as xr
+    from rasterio.transform import Affine
 
     if output is None:
         return array_to_memory_file(
-            array, source, dtype, compress, transpose, cellsize, crs, driver, **kwargs
+            array,
+            source,
+            dtype,
+            compress,
+            transpose,
+            cellsize,
+            crs=crs,
+            transform=transform,
+            driver=driver,
+            **kwargs,
         )
 
     if isinstance(array, xr.DataArray):
@@ -10447,7 +10466,7 @@ def array_to_image(
                 "crs must be provided if source is not provided, such as EPSG:3857"
             )
 
-        if "transform" not in kwargs:
+        if transform is None:
             # Define the geotransformation parameters
             xmin, ymin, xmax, ymax = (
                 0,
@@ -10458,8 +10477,12 @@ def array_to_image(
             transform = rasterio.transform.from_bounds(
                 xmin, ymin, xmax, ymax, array.shape[1], array.shape[0]
             )
-        else:
-            transform = kwargs["transform"]
+        elif isinstance(transform, Affine):
+            pass
+        elif isinstance(transform, (tuple, list)):
+            transform = Affine(*transform)
+
+        kwargs["transform"] = transform
 
     if dtype is None:
         # Determine the minimum and maximum values in the array
@@ -10509,6 +10532,7 @@ def array_to_image(
             for i in range(array.shape[2]):
                 dst.write(array[:, :, i], i + 1)
 
+
 def images_to_tiles(
     images: Union[str, List[str]], names: List[str] = None, **kwargs
 ) -> Dict[str, ipyleaflet.TileLayer]: