reinstate COG reader, make window boundless, fix dst/src window transform

Author: Lkruitwagen

src/satextractor/extractor/extractor.py

@@ -14,6 +14,39 @@
 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.
 
@@ -47,6 +80,68 @@ 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,
+                )
+
+        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,
@@ -128,7 +223,11 @@ def task_mosaic_patches(
     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]