diff --git a/config.toml b/config.toml index ad9e25d..bf631af 100644 --- a/config.toml +++ b/config.toml @@ -15,6 +15,10 @@ arcticdem-elevation-vrt = "data/input/ArcticDEM/elevation.vrt" arcticdem-dir = "data/download/arcticdem" cache-dir = "data/download" +[darts.preprocess] +tpi-outer-radius = 30 +tpi-inner-radius = 25 + [darts.segmentation] patch-size = 1024 overlap = 256 diff --git a/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py b/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py index 2e9d0ef..44a07a2 100644 --- a/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py +++ b/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py @@ -8,10 +8,10 @@ from typing import Literal import geopandas as gpd +import odc.geo import odc.geo.xr import pystac import requests -import shapely import xarray as xr import zarr import zarr.storage @@ -293,21 +293,21 @@ def procedural_download_datacube(storage: zarr.storage.Store, scenes: gpd.GeoDat def load_arcticdem_tile( - reference_dataset: xr.Dataset, + geobox: GeoBox, data_dir: Path, resolution: RESOLUTIONS, chunk_size: int = 6000, - buffer: int = 256, + buffer: int = 0, ) -> xr.Dataset: - """Get the corresponding ArcticDEM tile for the given reference dataset. + """Get the corresponding ArcticDEM tile for the given geobox. Args: - reference_dataset (xr.Dataset): The reference dataset. + geobox (GeoBox): The geobox for which the tile should be loaded. data_dir (Path): The directory where the ArcticDEM data is stored. resolution (Literal[2, 10, 32]): The resolution of the ArcticDEM data in m. chunk_size (int, optional): The chunk size for the datacube. Only relevant for the initial creation. Has no effect otherwise. Defaults to 6000. - buffer (int, optional): The buffer around the reference dataset in pixels. Defaults to 256. + buffer (int, optional): The buffer around the geobox in pixels. Defaults to 0. Returns: xr.Dataset: The ArcticDEM tile, with a buffer applied. @@ -316,7 +316,7 @@ def load_arcticdem_tile( Warning: 1. This function is not thread-safe. Thread-safety might be added in the future. - 2. Reference dataset must be in a meter based CRS. + 2. Geobox must be in a meter based CRS. """ # TODO: What is a good chunk size? @@ -330,9 +330,8 @@ def load_arcticdem_tile( storage = zarr.storage.FSStore(datacube_fpath) logger.debug(f"Getting ArcticDEM tile from {datacube_fpath.resolve()}") - # ! The reference dataset must be in a meter based CRS - reference_resolution = abs(int(reference_dataset.x[1] - reference_dataset.x[0])) - logger.debug(f"Found a reference resolution of {reference_resolution}m") + # ! The geobox must be in a meter based CRS + logger.debug(f"Found a reference resolution of {geobox.resolution.x}m") # Check if the zarr data already exists if not datacube_fpath.exists(): @@ -348,12 +347,9 @@ def load_arcticdem_tile( download_arcticdem_extent(data_dir) extent = gpd.read_parquet(extent_fpath) - # Add a buffer around the reference dataset to get the adjacent scenes - buffer_m = buffer * reference_resolution # nbuffer pixels * the resolution of the reference dataset - reference_bbox = shapely.geometry.box(*reference_dataset.rio.transform_bounds("epsg:3413")).buffer( - buffer_m, join_style="mitre" - ) - adjacent_scenes = extent[extent.intersects(reference_bbox)] + # Add a buffer around the geobox to get the adjacent scenes + reference_bbox = geobox.to_crs("epsg:3413", resolution=resolution).pad(buffer) + adjacent_scenes = extent[extent.intersects(reference_bbox.extent.geom)] # Download the adjacent scenes (if necessary) procedural_download_datacube(storage, adjacent_scenes) @@ -361,15 +357,8 @@ def load_arcticdem_tile( # Load the datacube and set the spatial_ref since it is set as a coordinate within the zarr format arcticdem_datacube = xr.open_zarr(storage, mask_and_scale=False).set_coords("spatial_ref") - # Get an AOI slice of the datacube, since rio.reproject_match would load the whole datacube - # Note: The bounds are not equal to the bbox orientation, because of the change of the CRS - xmin, ymin, xmax, ymax = reference_bbox.bounds - aoi_slice = { - "x": slice(xmin, xmax), - "y": slice(ymax, ymin), - } - logger.debug(f"AOI slice: {aoi_slice}") - arcticdem_aoi = arcticdem_datacube.sel(aoi_slice) + # Get an AOI slice of the datacube + arcticdem_aoi = arcticdem_datacube.odc.crop(reference_bbox.extent) # Change dtype of the datamask to uint8 for later reproject_match arcticdem_aoi["datamask"] = arcticdem_aoi.datamask.astype("uint8") diff --git a/darts-acquisition/src/darts_acquisition/s2.py b/darts-acquisition/src/darts_acquisition/s2.py index c2266d4..3f6a820 100644 --- a/darts-acquisition/src/darts_acquisition/s2.py +++ b/darts-acquisition/src/darts_acquisition/s2.py @@ -4,10 +4,10 @@ import time from pathlib import Path -import rasterio -import rasterio.enums +import odc.geo.xr # noqa: F401 import rioxarray # noqa: F401 import xarray as xr +from odc.geo.geobox import GeoBox logger = logging.getLogger(__name__.replace("darts_", "darts.")) @@ -61,12 +61,12 @@ def load_s2_scene(fpath: str | Path) -> xr.Dataset: return ds_s2 -def load_s2_masks(fpath: str | Path, reference_dataset: xr.Dataset) -> xr.Dataset: +def load_s2_masks(fpath: str | Path, reference_geobox: GeoBox) -> xr.Dataset: """Load the valid and quality data masks from a Sentinel 2 scene. Args: fpath (str | Path): The path to the directory containing the TIFF files. - reference_dataset (xr.Dataset): The reference dataset to reproject, resampled and cropped the masks data to. + reference_geobox (GeoBox): The reference geobox to reproject, resample and crop the masks data to. Raises: FileNotFoundError: If no matching TIFF file is found in the specified path. @@ -91,7 +91,7 @@ def load_s2_masks(fpath: str | Path, reference_dataset: xr.Dataset) -> xr.Datase # See scene classes here: https://custom-scripts.sentinel-hub.com/custom-scripts/sentinel-2/scene-classification/ da_scl = xr.open_dataarray(scl_path) - da_scl = da_scl.rio.reproject_match(reference_dataset, sampling=rasterio.enums.Resampling.nearest) + da_scl = da_scl.odc.reproject(reference_geobox, sampling="nearest") # valid data mask: valid data = 1, no data = 0 valid_data_mask = ( diff --git a/darts/src/darts/native.py b/darts/src/darts/native.py index 1d86651..f5b4ee3 100644 --- a/darts/src/darts/native.py +++ b/darts/src/darts/native.py @@ -246,7 +246,7 @@ def run_native_planet_pipeline_fast( # Find all PlanetScope orthotiles for fpath, outpath in planet_file_generator(orthotiles_dir, scenes_dir, output_data_dir): optical = load_planet_scene(fpath) - arcticdem = load_arcticdem_tile(optical, arcticdem_dir, resolution=2, buffer=tpi_outer_radius) + arcticdem = load_arcticdem_tile(optical.odc.geobox, arcticdem_dir, resolution=2, buffer=tpi_outer_radius) tcvis = load_tcvis(optical, cache_dir) data_masks = load_planet_masks(fpath) @@ -356,7 +356,7 @@ def run_native_sentinel2_pipeline( optical = load_s2_scene(fpath) arcticdem = load_arcticdem_from_vrt(arcticdem_slope_vrt, arcticdem_elevation_vrt, optical) tcvis = load_tcvis(optical, cache_dir) - data_masks = load_s2_masks(fpath, optical) + data_masks = load_s2_masks(fpath, optical.odc.geobox) tile = preprocess_legacy(optical, arcticdem, tcvis, data_masks)