geoai-py 0.1.7__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

geoai/__init__.py

@@ -2,7 +2,34 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "0.1.7"
+__version__ = "0.2.1"
 
 
+import os
+import sys
+
+
+def set_proj_lib_path():
+    """Set the PROJ_LIB environment variable based on the current conda environment."""
+    try:
+        # Get conda environment path
+        conda_env_path = os.environ.get("CONDA_PREFIX") or sys.prefix
+
+        # Set PROJ_LIB environment variable
+        proj_path = os.path.join(conda_env_path, "share", "proj")
+        gdal_path = os.path.join(conda_env_path, "share", "gdal")
+
+        # Check if the directory exists before setting
+        if os.path.exists(proj_path):
+            os.environ["PROJ_LIB"] = proj_path
+        if os.path.exists(gdal_path):
+            os.environ["GDAL_DATA"] = gdal_path
+    except Exception as e:
+        print(e)
+        return
+
+
+if "google.colab" not in sys.modules:
+    set_proj_lib_path()
+
 from .geoai import *

geoai/common.py

@@ -5,8 +5,12 @@ from collections.abc import Iterable
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, Callable
 import matplotlib.pyplot as plt
 
+import leafmap
 import torch
 import numpy as np
+import xarray as xr
+import rioxarray
+import rasterio as rio
 from torch.utils.data import DataLoader
 from torchgeo.datasets import RasterDataset, stack_samples, unbind_samples, utils
 from torchgeo.samplers import RandomGeoSampler, Units
@@ -55,10 +59,12 @@ def viz_raster(
     Returns:
         leafmap.Map: The map object with the raster layer added.
     """
-    import leafmap
 
     m = leafmap.Map(basemap=basemap)
 
+    if isinstance(source, dict):
+        source = dict_to_image(source)
+
     m.add_raster(
         source=source,
         indexes=indexes,
@@ -86,6 +92,7 @@ def viz_image(
     scale_factor: float = 1.0,
     figsize: Tuple[int, int] = (10, 5),
     axis_off: bool = True,
+    title: Optional[str] = None,
     **kwargs: Any,
 ) -> None:
     """
@@ -98,6 +105,7 @@ def viz_image(
         scale_factor (float, optional): The scale factor to apply to the image. Defaults to 1.0.
         figsize (Tuple[int, int], optional): The size of the figure. Defaults to (10, 5).
         axis_off (bool, optional): Whether to turn off the axis. Defaults to True.
+        title (Optional[str], optional): The title of the plot. Defaults to None.
         **kwargs (Any): Additional keyword arguments for plt.imshow().
 
     Returns:
@@ -124,6 +132,8 @@ def viz_image(
     plt.imshow(image, **kwargs)
     if axis_off:
         plt.axis("off")
+    if title is not None:
+        plt.title(title)
     plt.show()
     plt.close()
 
@@ -277,3 +287,150 @@ def calc_stats(
     # at the end, we shall have 2 vectors with length n=chnls
     # we will average them considering the number of images
     return accum_mean / len(files), accum_std / len(files)
+
+
+def dict_to_rioxarray(data_dict: Dict) -> xr.DataArray:
+    """Convert a dictionary to a xarray DataArray. The dictionary should contain the
+    following keys: "crs", "bounds", and "image". It can be generated from a TorchGeo
+    dataset sampler.
+
+    Args:
+        data_dict (Dict): The dictionary containing the data.
+
+    Returns:
+        xr.DataArray: The xarray DataArray.
+    """
+
+    from affine import Affine
+
+    # Extract components from the dictionary
+    crs = data_dict["crs"]
+    bounds = data_dict["bounds"]
+    image_tensor = data_dict["image"]
+
+    # Convert tensor to numpy array if needed
+    if hasattr(image_tensor, "numpy"):
+        # For PyTorch tensors
+        image_array = image_tensor.numpy()
+    else:
+        # If it's already a numpy array or similar
+        image_array = np.array(image_tensor)
+
+    # Calculate pixel resolution
+    width = image_array.shape[2]  # Width is the size of the last dimension
+    height = image_array.shape[1]  # Height is the size of the middle dimension
+
+    res_x = (bounds.maxx - bounds.minx) / width
+    res_y = (bounds.maxy - bounds.miny) / height
+
+    # Create the transform matrix
+    transform = Affine(res_x, 0.0, bounds.minx, 0.0, -res_y, bounds.maxy)
+
+    # Create dimensions
+    x_coords = np.linspace(bounds.minx + res_x / 2, bounds.maxx - res_x / 2, width)
+    y_coords = np.linspace(bounds.maxy - res_y / 2, bounds.miny + res_y / 2, height)
+
+    # If time dimension exists in the bounds
+    if hasattr(bounds, "mint") and hasattr(bounds, "maxt"):
+        # Create a single time value or range if needed
+        t_coords = [
+            bounds.mint
+        ]  # Or np.linspace(bounds.mint, bounds.maxt, num_time_steps)
+
+        # Create DataArray with time dimension
+        dims = (
+            ("band", "y", "x")
+            if image_array.shape[0] <= 10
+            else ("time", "band", "y", "x")
+        )
+
+        if dims[0] == "band":
+            # For multi-band single time
+            da = xr.DataArray(
+                image_array,
+                dims=dims,
+                coords={
+                    "band": np.arange(1, image_array.shape[0] + 1),
+                    "y": y_coords,
+                    "x": x_coords,
+                },
+            )
+        else:
+            # For multi-time multi-band
+            da = xr.DataArray(
+                image_array,
+                dims=dims,
+                coords={
+                    "time": t_coords,
+                    "band": np.arange(1, image_array.shape[1] + 1),
+                    "y": y_coords,
+                    "x": x_coords,
+                },
+            )
+    else:
+        # Create DataArray without time dimension
+        da = xr.DataArray(
+            image_array,
+            dims=("band", "y", "x"),
+            coords={
+                "band": np.arange(1, image_array.shape[0] + 1),
+                "y": y_coords,
+                "x": x_coords,
+            },
+        )
+
+    # Set spatial attributes
+    da.rio.write_crs(crs, inplace=True)
+    da.rio.write_transform(transform, inplace=True)
+
+    return da
+
+
+def dict_to_image(
+    data_dict: Dict[str, Any], output: Optional[str] = None, **kwargs
+) -> rio.DatasetReader:
+    """Convert a dictionary containing spatial data to a rasterio dataset or save it to
+    a file. The dictionary should contain the following keys: "crs", "bounds", and "image".
+    It can be generated from a TorchGeo dataset sampler.
+
+    This function transforms a dictionary with CRS, bounding box, and image data
+    into a rasterio DatasetReader using leafmap's array_to_image utility after
+    first converting to a rioxarray DataArray.
+
+    Args:
+        data_dict: A dictionary containing:
+            - 'crs': A pyproj CRS object
+            - 'bounds': A BoundingBox object with minx, maxx, miny, maxy attributes
+              and optionally mint, maxt for temporal bounds
+            - 'image': A tensor or array-like object with image data
+        output: Optional path to save the image to a file. If not provided, the image
+            will be returned as a rasterio DatasetReader object.
+        **kwargs: Additional keyword arguments to pass to leafmap.array_to_image.
+            Common options include:
+            - colormap: str, name of the colormap (e.g., 'viridis', 'terrain')
+            - vmin: float, minimum value for colormap scaling
+            - vmax: float, maximum value for colormap scaling
+
+    Returns:
+        A rasterio DatasetReader object that can be used for visualization or
+        further processing.
+
+    Examples:
+        >>> image = dict_to_image(
+        ...     {'crs': CRS.from_epsg(26911), 'bounds': bbox, 'image': tensor},
+        ...     colormap='terrain'
+        ... )
+        >>> fig, ax = plt.subplots(figsize=(10, 10))
+        >>> show(image, ax=ax)
+    """
+    da = dict_to_rioxarray(data_dict)
+
+    if output is not None:
+        out_dir = os.path.abspath(os.path.dirname(output))
+        if not os.path.exists(out_dir):
+            os.makedirs(out_dir, exist_ok=True)
+        da.rio.to_raster(output)
+        return output
+    else:
+        image = leafmap.array_to_image(da, **kwargs)
+        return image

geoai/download.py

@@ -26,6 +26,7 @@ def download_naip(
     output_dir: str,
     year: Optional[int] = None,
     max_items: int = 10,
+    overwrite: bool = False,
     preview: bool = False,
     **kwargs: Any,
 ) -> List[str]:
@@ -40,6 +41,7 @@ def download_naip(
         output_dir: Directory to save the downloaded imagery.
         year: Specific year of NAIP imagery to download (e.g., 2020). If None, returns imagery from all available years.
         max_items: Maximum number of items to download.
+        overwrite: If True, overwrite existing files with the same name.
         preview: If True, display a preview of the downloaded imagery.
 
     Returns:
@@ -75,6 +77,9 @@ def download_naip(
     search_results = catalog.search(**search_params)
     items = list(search_results.items())
 
+    if len(items) > max_items:
+        items = items[:max_items]
+
     if not items:
         print("No NAIP imagery found for the specified region and parameters.")
         return []
@@ -98,6 +103,10 @@ def download_naip(
             rgb_asset.href.split("?")[0]
         )  # Remove query parameters
         output_path = os.path.join(output_dir, original_filename)
+        if not overwrite and os.path.exists(output_path):
+            print(f"Skipping existing file: {output_path}")
+            downloaded_files.append(output_path)
+            continue
 
         print(f"Downloading item {i+1}/{len(items)}: {original_filename}")