Getpixels (#3)

* Added function to export faster * Adding conversion accepting np arrays * updated version * Add bounds_in_latlng to plot_segmentation_mask function --------- Co-authored-by: Gonzalo Mateo Garcia <[email protected]>
spaceml-org · Dec 4, 2023 · 8eb8194 · 8eb8194
1 parent 2ea2945
commit 8eb8194
Show file tree

Hide file tree

Showing 6 changed files with 156 additions and 41 deletions.
diff --git a/README.md b/README.md
@@ -5,7 +5,10 @@ Read data from rasters: very few dependencies, reads from cloud storage and lazy
 ## Install
 
 ```bash
-# Install with minimal requirements (only rasterio, numpy as shapely)
+# From pip
+pip install georeader-spaceml
+
+# From GitHub
 pip install git+https://github.com/spaceml-org/georeader#egg=georeader
 
 # Install with Google dependencies (to read objects from Google Cloud Storage or Google Earth Engine)

diff --git a/georeader/__init__.py b/georeader/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "1.0.5"
+__version__ = "1.0.6"
 
 import math
 from typing import Tuple, Any, Union

diff --git a/georeader/plot.py b/georeader/plot.py
@@ -202,7 +202,8 @@ def plot_segmentation_mask(mask:Union[GeoData, np.array], color_array:np.array,
                            interpretation_array:Optional[List[str]]=None,
                            legend:bool=True, ax:Optional[plt.Axes]=None,
                            add_scalebar:bool=False,
-                           kwargs_scalebar:Optional[dict]=None) -> plt.Axes:
+                           kwargs_scalebar:Optional[dict]=None,
+                           bounds_in_latlng:bool=True) -> plt.Axes:
     """
     Plots a discrete segmentation mask with a legend.
 
@@ -215,6 +216,7 @@ def plot_segmentation_mask(mask:Union[GeoData, np.array], color_array:np.array,
         add_scalebar (bool, optional): Defaults to False. Add a scalebar to the plot
         kwargs_scalebar (Optional[dict], optional): Defaults to None. Keyword arguments for the scalebar. 
         See https://github.com/ppinard/matplotlib-scalebar. (install with pip install matplotlib-scalebar)
+        bounds_in_latlng (bool, optional): Defaults to True. If True, the x and y ticks are shown in latlng.
     
     Returns:
         plt.Axes
@@ -268,6 +270,27 @@ def plot_segmentation_mask(mask:Union[GeoData, np.array], color_array:np.array,
             kwargs_scalebar["dx"] = 1
 
         ax.add_artist(ScaleBar(**kwargs_scalebar))
+
+    if bounds_in_latlng:
+        from matplotlib.ticker import FuncFormatter
+
+        xmin, ymin, xmax, ymax = mask.bounds
+
+        @FuncFormatter
+        def x_formatter(x, pos):
+            # transform x,ymin to latlng
+            longs, lats = rasterio.warp.transform(mask.crs, "epsg:4326", [x], [ymin])
+            return f"{longs[0]:.2f}"
+
+
+        @FuncFormatter
+        def y_formatter(y, pos):
+            # transform xmin,y to latlng
+            longs, lats = rasterio.warp.transform(mask.crs, "epsg:4326", [xmin], [y])
+            return f"{lats[0]:.2f}"
+
+        ax.xaxis.set_major_formatter(x_formatter)
+        ax.yaxis.set_major_formatter(y_formatter)
 
     return ax
 
diff --git a/georeader/readers/ee_image.py b/georeader/readers/ee_image.py
@@ -1,12 +1,19 @@
 from georeader.geotensor import GeoTensor, concatenate
 from georeader.abstract_reader import GeoData
 from georeader.rasterio_reader import RasterioReader
-from shapely.geometry import Polygon, MultiPolygon, mapping
-from typing import Union, Dict, Optional, Tuple, List
+from shapely.geometry import Polygon, MultiPolygon, mapping, box
+from typing import Union, Dict, Optional, Tuple, List, Any
 import ee
 import rasterio.windows
 from rasterio import Affine
 import numpy as np
+from collections import namedtuple
+from georeader import read, window_utils
+from io import BytesIO
+import rasterio
+from georeader import mosaic
+
+FakeGeoData=namedtuple("FakeGeoData",["crs", "transform"])
 
 
 def export_image_fast(image:ee.Image, geometry:Union[ee.Geometry, Polygon, MultiPolygon],
@@ -59,34 +66,94 @@ def export_image_fast(image:ee.Image, geometry:Union[ee.Geometry, Polygon, Multi
 
     return out
 
-def export_image_reproject(image: ee.Image,
-                           transform:rasterio.Affine,
-                           crs:str,
-                           dimensions:Tuple[int, int],
-                           bands:Optional[List[str]]=None,
-                           donwload_path:Optional[str]=None) -> GeoData:
-    # TODO default if not provided
-    donwload_path
-
-    download_url = image.getDownloadURL(params={
-        "name": name,
-        "bands": bands,
-        "crs_transform": list(transform)[:6],
-        "crs": str(crs).upper(),
-        "dimensions": dimensions,
-        "filePerBand": False})
-
-    import requests
-    r = requests.get(download_url, stream=True)
-    filenamezip = f'/home/gonzalo/Downloads/{name}.zip'
-    with open(filenamezip, "wb") as fd:
-        for chunk in r.iter_content(chunk_size=1024):
-            fd.write(chunk)
-
-    # TODO unzip?
-    filename = f'zip+file:///{filenamezip}!{name}.tif'
-    data = RasterioReader(filename)
-    return data
+def split_bounds(bounds:Tuple[float, float, float, float]) -> List[Tuple[float, float, float, float]]:
+    min_x, min_y, max_x, max_y = bounds
+    mid_x = (min_x + max_x) / 2
+    mid_y = (min_y + max_y) / 2
+
+    return [
+        (min_x, min_y, mid_x, mid_y),  # Lower left quadrant
+        (mid_x, min_y, max_x, mid_y),  # Lower right quadrant
+        (min_x, mid_y, mid_x, max_y),  # Upper left quadrant
+        (mid_x, mid_y, max_x, max_y),  # Upper right quadrant
+    ]
+
+def export_image_getpixels(asset_id: str,
+                           geometry:Union[Polygon, MultiPolygon],
+                           proj:Dict[str, Any],
+                           bands_gee:List[str],
+                           crs_polygon:str="EPSG:4326") -> GeoTensor:
+    """
+    Exports an image from the GEE as a GeoTensor. It uses the `ee.data.getPixels` method to export.
+
+    Args:
+        asset_id (str): Name of the asset
+        geometry (Union[Polygon, MultiPolygon]): geometry to export
+        proj (Dict[str, Any]): Dict with fields:
+            - crs: crs of the image
+            - transform: transform of the image
+        bands_gee (List[str]): List of bands to export
+        crs_polygon (str, optional): crs of the geometry. Defaults to "EPSG:4326".
+
+    Returns:
+        GeoTensor: GeoTensor object
+    """
+    geodata = FakeGeoData(crs=proj["crs"], transform=Affine(*proj["transform"]))
+    window_polygon = read.window_from_polygon(geodata, geometry, crs_polygon=crs_polygon,
+                                              window_surrounding=True)
+    window_polygon = window_utils.round_outer_window(window_polygon)
+    transform_window = rasterio.windows.transform(window_polygon, geodata.transform)
+
+    try:
+        data_raw = ee.data.getPixels({"assetId": asset_id, 
+                    'fileFormat':"GEO_TIFF", 
+                    'bandIds':  bands_gee,
+                    'grid': {
+                        'dimensions': {
+                            'height': window_polygon.height, 
+                            'width': window_polygon.width
+                        },
+                        'affineTransform': {
+                            'scaleX': transform_window.a,
+                            'shearX': transform_window.b,
+                            'translateX': transform_window.c,
+                            'shearY': transform_window.d,
+                            'scaleY': transform_window.e,
+                            'translateY': transform_window.f
+                        },
+                        'crsCode': geodata.crs
+                    }
+                    })
+        data = rasterio.open(BytesIO(data_raw))
+        geotensor = GeoTensor(data.read(), transform=data.transform,
+                             crs=data.crs, fill_value_default=data.nodata)
+    except ee.EEException as e:
+        # Check if the exception starts with Total request size
+        if str(e).startswith("Total request size"):
+            # Split the geometry in two and call recursively
+            bounds = geometry.bounds
+            # Split the bounds in two
+            geotensors = []
+            for sb in split_bounds(bounds):
+                # Create a polygon with the bounds
+                poly = box(*sb)
+                # Call recursively
+                gt = export_image_getpixels(asset_id, poly, 
+                                            proj, bands_gee, crs_polygon)
+                # Concatenate the GeoTensor
+                geotensors.append(gt)
+
+            dst_crs = geotensors[0].crs
+            aoi_dst_crs = window_utils.polygon_to_crs(geometry, 
+                                                      crs_polygon=crs_polygon, 
+                                                      dst_crs=dst_crs)
+
+            geotensor = mosaic.spatial_mosaic(geotensors, 
+                                              polygon=aoi_dst_crs, 
+                                              dst_crs=dst_crs)
+        else:
+            raise e
+    return geotensor
 
 
 

diff --git a/georeader/readers/ee_query.py b/georeader/readers/ee_query.py
@@ -97,7 +97,7 @@ def query(area:Union[MultiPolygon,Polygon],
           return_collection:bool=False,
           add_s2cloudless:bool=False)-> Union[gpd.GeoDataFrame, Tuple[gpd.GeoDataFrame, ee.ImageCollection]]:
     """
-    Query Landsat and Sentinel-2 products from the Google Earth Engine
+    Query Landsat and Sentinel-2 products from the Google Earth Engine.
 
     Args:
         area: area to query images in EPSG:4326
@@ -227,6 +227,10 @@ def query(area:Union[MultiPolygon,Polygon],
 
     geodf = _add_stuff(geodf, area, tz)
 
+    # Fix ids of Landsat to remove initial shit in the names
+    if geodf.satellite.str.startswith("LC0").any():
+        geodf.loc[geodf.satellite.str.startswith("LC0"),"gee_id"] = geodf.loc[geodf.satellite.str.startswith("LC0"),"gee_id"].apply(lambda x: "LC0"+x.split("LC0")[1])
+
     if filter_duplicates:
         # TODO filter prioritizing s2cloudless?
         geodf = query_utils.filter_products_overlap(area, geodf,

diff --git a/georeader/reflectance.py b/georeader/reflectance.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pandas as pd
 import pkg_resources
+from numpy.typing import ArrayLike, NDArray
 
 
 def earth_sun_distance_correction_factor(date_of_acquisition:datetime) -> float:
@@ -30,8 +31,9 @@ def earth_sun_distance_correction_factor(date_of_acquisition:datetime) -> float:
     return 1 - 0.01673 * np.cos(0.0172 * (tm_yday - 4))
 
 
-def observation_date_correction_factor(center_coords:Tuple[float, float], date_of_acquisition:datetime,
-                                       crs_coords:Optional[str]=None,) -> float:
+def observation_date_correction_factor(center_coords:Tuple[float, float], 
+                                       date_of_acquisition:datetime,
+                                       crs_coords:Optional[str]=None) -> float:
     """
     returns  (pi * d^2) / cos(solarzenithangle/180*pi)
 
@@ -65,8 +67,10 @@ def observation_date_correction_factor(center_coords:Tuple[float, float], date_o
     return np.pi*(d**2) / np.cos(sza/180.*np.pi)
 
 
-def radiance_to_reflectance(data:GeoTensor, solar_irradiance:Union[List[float], np.array],
-                            date_of_acquisition:datetime) -> GeoTensor:
+def radiance_to_reflectance(data:Union[GeoTensor, ArrayLike], solar_irradiance:Union[List[float], np.array],
+                            date_of_acquisition:datetime,
+                            center_coords:Optional[Tuple[float, float]]=None,
+                            crs_coords:Optional[str]=None ) -> Union[GeoTensor, NDArray]:
     """
 
     toaBandX = (radianceBandX / 100 * pi * d^2) / (cos(solarzenithangle/180*pi) * solarIrradianceBandX)
@@ -83,6 +87,8 @@ def radiance_to_reflectance(data:GeoTensor, solar_irradiance:Union[List[float],
                 microwatts per centimeter_squared per nanometer per steradian
         solar_irradiance: (C,) vector units: W/m²/nm
         date_of_acquisition: date of acquisition to compute the solar zenith angles
+        center_coords: location being considered (x,y) (long, lat if EPSG:4326)
+        crs_coords: if None it will assume center_coords are in EPSG:4326
 
     Returns:
         GeoTensor with ToA on each channel
@@ -94,14 +100,26 @@ def radiance_to_reflectance(data:GeoTensor, solar_irradiance:Union[List[float],
         f"Different number of channels {data.shape[0]} than number of radiances {len(solar_irradiance)}"
 
     # Get latitude and longitude of the center of image to compute the solar angle
-    center_coords = data.transform * (data.shape[-1] // 2, data.shape[-2] // 2)
-    constant_factor = observation_date_correction_factor(center_coords, date_of_acquisition, crs_coords=data.crs)
+    if center_coords is None:
+        assert isinstance(data, GeoTensor), "If center_coords is None, data must be a GeoTensor"
+        center_coords = data.transform * (data.shape[-1] // 2, data.shape[-2] // 2)
+        crs_coords = data.crs
+
+    constant_factor = observation_date_correction_factor(center_coords, date_of_acquisition, crs_coords=crs_coords)
+
+    if isinstance(data, GeoTensor):
+        data_values = data.values
+    else:
+        data_values = data
 
     # µW /(nm cm² sr) to W/(nm m² sr)
-    radiances = data.values * (10**(-6) / 1) * (1 /10**(-4))
+    radiances = data_values * (10**(-6) / 1) * (1 /10**(-4))
 
     # data_toa = data.values / 100 * constant_factor / solar_irradiance
     data_toa = radiances * constant_factor / solar_irradiance
+    if not  isinstance(data, GeoTensor):
+        return data_toa
+
     mask = data.values == data.fill_value_default
     data_toa[mask] = data.fill_value_default