Merge pull request #27 from FrontierDevelopmentLab/bugfix/boundless-window

fix COG reader

Author: frandorr

providers/gcp/main.py

@@ -157,7 +157,7 @@ def extract_patches():
             cloud_fs=fs,
             download_f=download_blob,
             task=task,
-            method="first",
+            method="max",
             resolution=archive_resolution,
         )
 

src/satextractor/extractor/extractor.py

@@ -6,14 +6,51 @@
 
 import numpy as np
 import rasterio
+from affine import Affine
 from loguru import logger
 from osgeo import gdal
 from osgeo import osr
+from rasterio import warp
+from rasterio.crs import CRS
+from rasterio.enums import Resampling
 from rasterio.merge import merge as riomerge
 from satextractor.models import ExtractionTask
 from satextractor.models import Tile
 
 
+def get_window_union(
+    tiles: List[Tile],
+    ds: rasterio.io.DatasetReader,
+) -> rasterio.windows.Window:
+
+    """Get the window union to read all tiles from the geotiff.
+
+    Args:
+        tiles (List[Tile]): the tiles
+        ds (rasterio.io.DatasetReader): the rasterio dataset to read (for the transform)
+
+    Returns:
+        rasterio.windows.Window: The union of all tile windows.
+    """
+
+    windows = []
+
+    for tile in tiles:
+
+        bounds_arr_tile_crs = tile.bbox
+        bounds_arr_rast_crs = warp.transform_bounds(
+            CRS.from_epsg(tile.epsg),
+            ds.crs,
+            *bounds_arr_tile_crs,
+        )
+
+        window = rasterio.windows.from_bounds(*bounds_arr_rast_crs, ds.transform)
+
+        windows.append(window)
+
+    return rasterio.windows.union(windows)
+
+
 def get_proj_win(tiles: List[Tile]) -> Tuple[int, int, int, int]:
     """Get the projection bounds window of the tiles.
 
@@ -48,6 +85,69 @@ def get_tile_pixel_coords(tiles: List[Tile], raster_file: str) -> List[Tuple[int
     return list(zip(rows, cols))
 
 
+def download_and_extract_tiles_window_COG(
+    fs: Any,
+    task: ExtractionTask,
+    resolution: int,
+) -> List[str]:
+    """Download and extract from the task assets the data for the window from each asset.
+
+    Args:
+        task (ExtractionTask): The extraction task
+        resolution (int): The target resolution
+
+    Returns:
+        List[str]: A list of files that store the crops of the original assets
+    """
+
+    # task tiles all have same CRS, so get their max extents and crs
+    left, top, right, bottom = get_proj_win(task.tiles)
+    epsg = task.tiles[0].epsg
+
+    # set the transforms for the output file
+    dst_transform = Affine(resolution, 0.0, left, 0.0, -resolution, top)
+    out_shp = (int((right - left) / resolution), int((top - bottom) / resolution))
+
+    outfiles = []
+
+    band = task.band
+    urls = [item.assets[band].href for item in task.item_collection.items]
+
+    for ii, url in enumerate(urls):
+        with fs.open(url) as f:
+            with rasterio.open(f) as ds:
+                window = get_window_union(task.tiles, ds)
+
+                rst_arr = ds.read(
+                    1,
+                    window=window,
+                    out_shape=out_shp,
+                    fill_value=0,
+                    boundless=True,
+                    resampling=Resampling.bilinear,
+                )
+
+        out_f = f"{task.task_id}_{ii}.tif"
+
+        with rasterio.open(
+            out_f,
+            "w",
+            driver="GTiff",
+            count=1,
+            width=out_shp[0],
+            height=out_shp[1],
+            transform=dst_transform,
+            crs=CRS.from_epsg(epsg),
+            dtype=rst_arr.dtype,
+        ) as dst:
+
+            dst.write(rst_arr, indexes=1)
+
+        outfiles.append(out_f)
+
+    return outfiles
+
+
 def download_and_extract_tiles_window(
     download_f: Callable,
     task: ExtractionTask,
@@ -112,7 +212,7 @@ def task_mosaic_patches(
     cloud_fs: Any,
     download_f: Callable,
     task: ExtractionTask,
-    method: str = "first",
+    method: str = "max",
     resolution: int = 10,
     dst_path="merged.jp2",
 ) -> List[np.ndarray]:
@@ -121,14 +221,18 @@ def task_mosaic_patches(
     Args:
         download_f (Callable): The function to download the task assets
         task (ExtractionTask): The task
-        method (str, optional): The method to use while merging the assets. Defaults to "first".
+        method (str, optional): The method to use while merging the assets. Defaults to "max".
         resolution (int, optional): The target resolution. Defaults to 10.
         dst_path (str): path to store the merged files
 
     Returns:
         List[np.ndarray]: The tile patches as numpy arrays
     """
-    out_files = download_and_extract_tiles_window(download_f, task, resolution)
+
+    if task.constellation == "sentinel-2":
+        out_files = download_and_extract_tiles_window(download_f, task, resolution)
+    else:
+        out_files = download_and_extract_tiles_window_COG(cloud_fs, task, resolution)
 
     out_f = f"{task.task_id}_{dst_path}"
     datasets = [rasterio.open(f) for f in out_files]