Building Sentinel-2 mosaics with RasterFlow

import os
from datetime import datetime

import geopandas as gpd
import wkls
import xarray as xr
from rasterflow_remote import RasterflowClient
from rasterflow_remote.data_models import DatasetEnum

rf_client = RasterflowClient()

# Generate a geometry for Haskell County, Kansas using Well-Known Locations (https://github.com/wherobots/wkls)
gdf = gpd.read_file(wkls["us"]["ks"]["Haskell County"].geojson())

# Save the geometry to a parquet file in the user's S3 path
aoi_path = os.getenv("USER_S3_PATH") + "harvest_mosaic_aoi.parquet"
gdf.to_parquet(aoi_path)

# Store bounds for visualization later
min_lon, min_lat, max_lon, max_lat = gdf.total_bounds

print(f"AOI saved to: {aoi_path}")
print(f"Bounds: [{min_lon:.4f}, {min_lat:.4f}, {max_lon:.4f}, {max_lat:.4f}]")

# Use the pre-generated mosaic instead
# mosaic_store = "s3://wherobots-examples/rasterflow/mosaics/haskell_harvest.zarr"

mosaic_index = rf_client.build_mosaic(
    # Dataset type for harvest season median mosaic
    datasets=[DatasetEnum.S2_MED_HARVEST],
    
    # Path to our AOI in GeoParquet format
    aoi=aoi_path,
    
    # Date range for imagery to be used
    start=datetime(2024, 1, 1),
    end=datetime(2025, 1, 1),
    
    # Output CRS (Web Mercator)
    crs_epsg=3857,
)

mosaic_index.mosaics

# This example AOI has one geometry, so we select the first (only) mosaic location.
mosaic_store = mosaic_index.first_row_mosaic

mosaic_store

all_bands_index = rf_client.build_mosaic(
    # All 12 L2A spectral bands, median composite over the date window below
    datasets=[DatasetEnum.S2_MED_WINDOWED_PIXEL_ALL_BANDS],
    
    # Reuse the same AOI as the harvest example
    aoi=aoi_path,
    
    # Custom date window — used directly, no seasonal heuristic
    start=datetime(2024, 7, 1),
    end=datetime(2024, 9, 1),
    
    # Output CRS (Web Mercator)
    crs_epsg=3857,
)

all_bands_store = all_bands_index.first_row_mosaic
print(f"All-bands mosaic store: {all_bands_store}")