Merge pull request #26 from FrontierDevelopmentLab/bugfix/corrupted-bands

Remove COG downloader and use same as jp2k because it was buggy

Author: frandorr

src/satextractor/extractor/extractor.py

@@ -6,51 +6,14 @@
 
 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(
-            ds.crs,
-            CRS.from_epsg(tile.epsg),
-            *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.
 
@@ -85,69 +48,6 @@ 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))
-    logger.debug(f"Affine with resolution: {dst_transform} and out_shp {out_shp} ")
-
-    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,
-            resampling=Resampling.bilinear,
-        ) 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,
@@ -228,11 +128,7 @@ def task_mosaic_patches(
     Returns:
         List[np.ndarray]: The tile patches as numpy arrays
     """
-
-    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_files = download_and_extract_tiles_window(download_f, task, resolution)
 
     out_f = f"{task.task_id}_{dst_path}"
     datasets = [rasterio.open(f) for f in out_files]