rslearn.data_sources.aws_open_data

aws_open_data

Data source for raster data in Registry of Open Data on AWS.

NaipItem

Bases: Item

An item in the Naip data source.

Source code in rslearn/data_sources/aws_open_data.py

class NaipItem(Item):
    """An item in the Naip data source."""

    def __init__(self, name: str, geometry: STGeometry, blob_path: str):
        """Creates a new NaipItem.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
            blob_path: path in bucket
        """
        super().__init__(name, geometry)
        self.blob_path = blob_path

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["blob_path"] = self.blob_path
        return d

    @staticmethod
    def deserialize(d: dict) -> "NaipItem":
        """Deserializes an item from a JSON-decoded dictionary."""
        item = super(NaipItem, NaipItem).deserialize(d)
        return NaipItem(
            name=item.name, geometry=item.geometry, blob_path=d["blob_path"]
        )

serialize

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/aws_open_data.py

def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["blob_path"] = self.blob_path
    return d

deserialize staticmethod

deserialize(d: dict) -> NaipItem

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/aws_open_data.py

@staticmethod
def deserialize(d: dict) -> "NaipItem":
    """Deserializes an item from a JSON-decoded dictionary."""
    item = super(NaipItem, NaipItem).deserialize(d)
    return NaipItem(
        name=item.name, geometry=item.geometry, blob_path=d["blob_path"]
    )

Naip

Bases: DataSource

A data source for NAIP imagery on AWS.

Specifically uses the naip-source requester pays bucket maintained by Esri. See https://registry.opendata.aws/naip/ for more information.

AWS credentials must be configured for use with boto3.

Source code in rslearn/data_sources/aws_open_data.py

class Naip(DataSource):
    """A data source for NAIP imagery on AWS.

    Specifically uses the naip-source requester pays bucket maintained by Esri. See
    https://registry.opendata.aws/naip/ for more information.

    AWS credentials must be configured for use with boto3.
    """

    bucket_name = "naip-source"

    manifest_fname = "manifest.txt"

    def __init__(
        self,
        index_cache_dir: str,
        use_rtree_index: bool = False,
        states: list[str] | None = None,
        years: list[int] | None = None,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new Naip instance.

        Args:
            index_cache_dir: directory to cache index shapefiles.
            use_rtree_index: whether to create an rtree index to enable faster lookups
                (default false)
            states: optional list of states (lowercase two-letter codes) to restrict
                the search. If use_rtree_index is enabled, the rtree will only be
                populated with data from these states.
            years: optional list of years to restrict the search
            context: the data source context.
        """
        # If context is provided, we join the directory with the dataset path,
        # otherwise we treat it directly as UPath.
        if context.ds_path is not None:
            self.index_cache_dir = join_upath(context.ds_path, index_cache_dir)
        else:
            self.index_cache_dir = UPath(index_cache_dir)

        self.states = states
        self.years = years

        self.index_cache_dir.mkdir(parents=True, exist_ok=True)

        self.bucket = boto3.resource("s3").Bucket(self.bucket_name)
        self.rtree_index: Any | None = None
        if use_rtree_index:
            from rslearn.utils.rtree_index import RtreeIndex, get_cached_rtree

            def build_fn(index: RtreeIndex) -> None:
                for item in self._read_index_shapefiles(desc="Building rtree index"):
                    index.insert(item.geometry.shp.bounds, json.dumps(item.serialize()))

            self.rtree_index = get_cached_rtree(self.index_cache_dir, build_fn)

    def _download_manifest(self) -> UPath:
        """Download the manifest that enumerates files in the bucket.

        Returns:
            The local filename where the manifest has been downloaded.
        """
        manifest_path = self.index_cache_dir / self.manifest_fname
        if not manifest_path.exists():
            with open_atomic(manifest_path, "wb") as dst:
                self.bucket.download_fileobj(
                    self.manifest_fname,
                    dst,
                    ExtraArgs={"RequestPayer": "requester"},
                )
        return manifest_path

    def _download_index_shapefiles(self) -> None:
        """Download all index shapefiles that specify image extents."""
        manifest_path = self._download_manifest()
        needed_files: list[tuple[str, UPath]] = []
        with manifest_path.open() as f:
            for line in f:
                blob_path = line.strip()
                if not blob_path:
                    continue
                if "/index/" not in blob_path:
                    continue

                # Data before 2012 doesn't seem to be present even though index files
                # might exist?
                # Also filter state/year here.
                path_parts = blob_path.split("/")
                state_code = path_parts[0]
                year = int(path_parts[1])
                if year < 2012:
                    continue
                if self.states and state_code not in self.states:
                    continue
                if self.years and year not in self.years:
                    continue

                cache_path = self.index_cache_dir / blob_path
                if cache_path.exists():
                    continue
                needed_files.append((blob_path, cache_path))
        for blob_path, cache_path in tqdm.tqdm(
            needed_files, desc="Downloading index files"
        ):
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            with open_atomic(cache_path, "wb") as dst:
                self.bucket.download_fileobj(
                    blob_path, dst, ExtraArgs={"RequestPayer": "requester"}
                )

    def _read_index_shapefiles(
        self, desc: str | None = None
    ) -> Generator[NaipItem, None, None]:
        """Read the index shapefiles and yield NaipItems corresponding to each image."""
        self._download_index_shapefiles()

        # Sometimes we see files like m_3410458_sw_13_060_20180612_20190426.tif but
        # the actual name is m_3410458_sw_13_060_20180612_20190427.tif (last part, the
        # version date, is different).
        # Also sometimes it is just m_4212362_sw_10_1_20140622.tif (ending in capture
        # date).
        # So here we create dict from just the prefix (i.e. exclude the version date).
        tif_blob_path_dict = {}
        manifest_path = self._download_manifest()
        with manifest_path.open() as f:
            for line in f:
                blob_path = line.strip()
                if not blob_path:
                    continue
                if not blob_path.endswith(".tif"):
                    continue
                parts = blob_path.split(".tif")[0].split("_")
                prefix = "_".join(parts[0:6])
                tif_blob_path_dict[prefix] = blob_path

        shape_files = self.index_cache_dir.glob("**/*.shp")
        if desc:
            shape_files = tqdm.tqdm(shape_files, desc=desc)
        for shp_fname in shape_files:
            # We copy the file to local since it needs to read all the auxiliary files
            # of the shapefile.
            prefix = ".".join(shp_fname.name.split(".")[:-1])
            aux_files: list[UPath] = []
            for ext in SHAPEFILE_AUX_EXTENSIONS:
                aux_files.append(shp_fname.parent / (prefix + ext))

            with get_upath_local(shp_fname, extra_paths=aux_files) as local_fname:
                with fiona.open(local_fname) as f:
                    src_crs = f.crs
                    dst_crs = fiona.crs.from_epsg(WGS84_EPSG)

                    for feature in f:
                        geometry = fiona.transform.transform_geom(
                            src_crs, dst_crs, feature["geometry"]
                        )
                        shp = shapely.geometry.shape(geometry)

                        # Properties specifies TIF filename like:
                        # - m_4212362_sw_10_1_20140622_20140923.tif.
                        # Index fname is like:
                        # - ca/2014/100cm/index/naip_3_14_3_1_ca.shp.
                        # So use that to reconstruct the prefix:
                        # - ca/2014/100cm/index/rgbir/m_4212362_sw_10_1_20140622
                        base_dir = "/".join(
                            shp_fname.parent.parent.path.split("/")[-3:]
                        )
                        tif_fname = feature["properties"]["FileName"]
                        fname_parts = tif_fname.split(".tif")[0].split("_")
                        tile_id = fname_parts[1][0:5]
                        fname_prefix = "_".join(fname_parts[:-1])
                        full_prefix = f"{base_dir}/rgbir/{tile_id}/{fname_prefix}"

                        if full_prefix not in tif_blob_path_dict:
                            print(
                                f"warning: skipping file {tif_fname} seen in "
                                + f"shapefile {shp_fname} but {full_prefix} does not exist "
                                + "in manifest"
                            )
                            continue
                        blob_path = tif_blob_path_dict[full_prefix]

                        # SrcImgDate is either string like "20180905" or int 20180905.
                        # We make sure it is string here.
                        # But it also could not exist in which case we need to fallback to
                        # extracting it from filename.
                        if "SrcImgDate" in feature["properties"]:
                            src_img_date = feature["properties"]["SrcImgDate"]
                            if isinstance(src_img_date, int):
                                src_img_date = str(src_img_date)
                        else:
                            src_img_date = fname_parts[5]
                        time = datetime.strptime(src_img_date, "%Y%m%d").replace(
                            tzinfo=UTC
                        )

                        geometry = STGeometry(WGS84_PROJECTION, shp, (time, time))

                        yield NaipItem(
                            name=blob_path.split("/")[-1].split(".tif")[0],
                            geometry=geometry,
                            blob_path=blob_path,
                        )

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[NaipItem]]]:
        """Get a list of items in the data source intersecting the given geometries.

        Args:
            geometries: the spatiotemporal geometries
            query_config: the query configuration

        Returns:
            List of groups of items that should be retrieved for each geometry.
        """
        wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]

        items: list = [[] for _ in geometries]
        if self.rtree_index:
            for idx, geometry in enumerate(wgs84_geometries):
                encoded_items: set = set()
                for shp in flatten_shape(geometry.shp):
                    encoded_items.update(self.rtree_index.query(shp.bounds))
                for encoded_item in encoded_items:
                    item = NaipItem.deserialize(json.loads(encoded_item))
                    if not item.geometry.shp.intersects(geometry.shp):
                        continue
                    items[idx].append(item)
        else:
            index = GridIndex(0.01)
            for idx, geometry in enumerate(wgs84_geometries):
                for shp in flatten_shape(geometry.shp):
                    index.insert(shp.bounds, idx)
            for item in self._read_index_shapefiles():
                results = index.query(item.geometry.shp.bounds)
                for idx in results:
                    geometry = wgs84_geometries[idx]
                    if not geometry.shp.intersects(item.geometry.shp):
                        continue
                    items[idx].append(item)

        groups = []
        for geometry, item_list in zip(wgs84_geometries, items):
            cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
                geometry, item_list, query_config
            )
            groups.append(cur_groups)
        return groups

    def deserialize_item(self, serialized_item: dict) -> NaipItem:
        """Deserializes an item from JSON-decoded data."""
        return NaipItem.deserialize(serialized_item)

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[NaipItem],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        Args:
            tile_store: the tile store to ingest into
            items: the items to ingest
            geometries: a list of geometries needed for each item
        """
        for item in items:
            bands = ["R", "G", "B", "IR"]
            if tile_store.is_raster_ready(item, bands):
                continue

            with tempfile.TemporaryDirectory() as tmp_dir:
                fname = os.path.join(tmp_dir, item.blob_path.split("/")[-1])
                self.bucket.download_file(
                    item.blob_path, fname, ExtraArgs={"RequestPayer": "requester"}
                )
                tile_store.write_raster_file(
                    item, bands, UPath(fname), time_range=item.geometry.time_range
                )

get_items

get_items(geometries: list[STGeometry], query_config: QueryConfig) -> list[list[MatchedItemGroup[NaipItem]]]

Get a list of items in the data source intersecting the given geometries.

Parameters:

Name	Type	Description	Default
geometries	list[STGeometry]	the spatiotemporal geometries	required
query_config	QueryConfig	the query configuration	required

Returns:

Type	Description
list[list[MatchedItemGroup[NaipItem]]]	List of groups of items that should be retrieved for each geometry.

Source code in rslearn/data_sources/aws_open_data.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[NaipItem]]]:
    """Get a list of items in the data source intersecting the given geometries.

    Args:
        geometries: the spatiotemporal geometries
        query_config: the query configuration

    Returns:
        List of groups of items that should be retrieved for each geometry.
    """
    wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]

    items: list = [[] for _ in geometries]
    if self.rtree_index:
        for idx, geometry in enumerate(wgs84_geometries):
            encoded_items: set = set()
            for shp in flatten_shape(geometry.shp):
                encoded_items.update(self.rtree_index.query(shp.bounds))
            for encoded_item in encoded_items:
                item = NaipItem.deserialize(json.loads(encoded_item))
                if not item.geometry.shp.intersects(geometry.shp):
                    continue
                items[idx].append(item)
    else:
        index = GridIndex(0.01)
        for idx, geometry in enumerate(wgs84_geometries):
            for shp in flatten_shape(geometry.shp):
                index.insert(shp.bounds, idx)
        for item in self._read_index_shapefiles():
            results = index.query(item.geometry.shp.bounds)
            for idx in results:
                geometry = wgs84_geometries[idx]
                if not geometry.shp.intersects(item.geometry.shp):
                    continue
                items[idx].append(item)

    groups = []
    for geometry, item_list in zip(wgs84_geometries, items):
        cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
            geometry, item_list, query_config
        )
        groups.append(cur_groups)
    return groups

deserialize_item

deserialize_item(serialized_item: dict) -> NaipItem

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/aws_open_data.py

def deserialize_item(self, serialized_item: dict) -> NaipItem:
    """Deserializes an item from JSON-decoded data."""
    return NaipItem.deserialize(serialized_item)

ingest

ingest(tile_store: TileStoreWithLayer, items: list[NaipItem], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Parameters:

Name	Type	Description	Default
tile_store	TileStoreWithLayer	the tile store to ingest into	required
items	list[NaipItem]	the items to ingest	required
geometries	list[list[STGeometry]]	a list of geometries needed for each item	required

Source code in rslearn/data_sources/aws_open_data.py

def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[NaipItem],
    geometries: list[list[STGeometry]],
) -> None:
    """Ingest items into the given tile store.

    Args:
        tile_store: the tile store to ingest into
        items: the items to ingest
        geometries: a list of geometries needed for each item
    """
    for item in items:
        bands = ["R", "G", "B", "IR"]
        if tile_store.is_raster_ready(item, bands):
            continue

        with tempfile.TemporaryDirectory() as tmp_dir:
            fname = os.path.join(tmp_dir, item.blob_path.split("/")[-1])
            self.bucket.download_file(
                item.blob_path, fname, ExtraArgs={"RequestPayer": "requester"}
            )
            tile_store.write_raster_file(
                item, bands, UPath(fname), time_range=item.geometry.time_range
            )

Sentinel2Modality

Bases: Enum

Sentinel-2 modality, either L1C or L2A.

Source code in rslearn/data_sources/aws_open_data.py

class Sentinel2Modality(Enum):
    """Sentinel-2 modality, either L1C or L2A."""

    L1C = "L1C"
    L2A = "L2A"

Sentinel2Item

Bases: Item

An item in the Sentinel2 data source.

Source code in rslearn/data_sources/aws_open_data.py

class Sentinel2Item(Item):
    """An item in the Sentinel2 data source."""

    def __init__(
        self, name: str, geometry: STGeometry, blob_path: str, cloud_cover: float
    ):
        """Creates a new Sentinel2Item.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
            blob_path: path in bucket, e.g. tiles/51/C/WM/2024/2/1/0/
            cloud_cover: the scene's cloud cover
        """
        super().__init__(name, geometry)
        self.blob_path = blob_path
        self.cloud_cover = cloud_cover

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["blob_path"] = self.blob_path
        d["cloud_cover"] = self.cloud_cover
        return d

    @staticmethod
    def deserialize(d: dict) -> "Sentinel2Item":
        """Deserializes an item from a JSON-decoded dictionary."""
        if "name" not in d:
            d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
        item = super(Sentinel2Item, Sentinel2Item).deserialize(d)
        return Sentinel2Item(
            name=item.name,
            geometry=item.geometry,
            blob_path=d["blob_path"],
            cloud_cover=d["cloud_cover"],
        )

serialize

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/aws_open_data.py

def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["blob_path"] = self.blob_path
    d["cloud_cover"] = self.cloud_cover
    return d

deserialize staticmethod

deserialize(d: dict) -> Sentinel2Item

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/aws_open_data.py

@staticmethod
def deserialize(d: dict) -> "Sentinel2Item":
    """Deserializes an item from a JSON-decoded dictionary."""
    if "name" not in d:
        d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
    item = super(Sentinel2Item, Sentinel2Item).deserialize(d)
    return Sentinel2Item(
        name=item.name,
        geometry=item.geometry,
        blob_path=d["blob_path"],
        cloud_cover=d["cloud_cover"],
    )

Sentinel2

Bases: ItemLookupDataSource[Sentinel2Item], RetrieveItemDataSource[Sentinel2Item]

A data source for Sentinel-2 L1C and L2A imagery on AWS.

Specifically, uses the sentinel-s2-l1c and sentinel-s2-l2a S3 buckets maintained by Sinergise. They state the data is "added regularly, usually within few hours after they are available on Copernicus OpenHub".

See https://aws.amazon.com/marketplace/pp/prodview-2ostsvrguftb2 for details about the buckets.

The buckets were previously requester pays but now appear to allow anonymous/free access, so AWS credentials are not needed.

Source code in rslearn/data_sources/aws_open_data.py

class Sentinel2(
    ItemLookupDataSource[Sentinel2Item], RetrieveItemDataSource[Sentinel2Item]
):
    """A data source for Sentinel-2 L1C and L2A imagery on AWS.

    Specifically, uses the sentinel-s2-l1c and sentinel-s2-l2a S3 buckets maintained by
    Sinergise. They state the data is "added regularly, usually within few hours after
    they are available on Copernicus OpenHub".

    See https://aws.amazon.com/marketplace/pp/prodview-2ostsvrguftb2 for details about
    the buckets.

    The buckets were previously requester pays but now appear to allow anonymous/free
    access, so AWS credentials are not needed.
    """

    bucket_names = {
        Sentinel2Modality.L1C: "sentinel-s2-l1c",
        Sentinel2Modality.L2A: "sentinel-s2-l2a",
    }

    band_fnames = {
        Sentinel2Modality.L1C: [
            ("B01.jp2", ["B01"]),
            ("B02.jp2", ["B02"]),
            ("B03.jp2", ["B03"]),
            ("B04.jp2", ["B04"]),
            ("B05.jp2", ["B05"]),
            ("B06.jp2", ["B06"]),
            ("B07.jp2", ["B07"]),
            ("B08.jp2", ["B08"]),
            ("B09.jp2", ["B09"]),
            ("B10.jp2", ["B10"]),
            ("B11.jp2", ["B11"]),
            ("B12.jp2", ["B12"]),
            ("B8A.jp2", ["B8A"]),
            ("TCI.jp2", ["R", "G", "B"]),
        ],
        Sentinel2Modality.L2A: [
            ("R20m/B01.jp2", ["B01"]),
            ("R10m/B02.jp2", ["B02"]),
            ("R10m/B03.jp2", ["B03"]),
            ("R10m/B04.jp2", ["B04"]),
            ("R20m/B05.jp2", ["B05"]),
            ("R20m/B06.jp2", ["B06"]),
            ("R20m/B07.jp2", ["B07"]),
            ("R10m/B08.jp2", ["B08"]),
            ("R60m/B09.jp2", ["B09"]),
            ("R20m/B11.jp2", ["B11"]),
            ("R20m/B12.jp2", ["B12"]),
            ("R20m/B8A.jp2", ["B8A"]),
            ("R10m/TCI.jp2", ["R", "G", "B"]),
        ],
    }

    def __init__(
        self,
        modality: Sentinel2Modality,
        metadata_cache_dir: str,
        sort_by: str | None = None,
        harmonize: bool = False,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new Sentinel2 instance.

        Args:
            modality: L1C or L2A.
            metadata_cache_dir: directory to cache product metadata files.
            sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                SpaceMode.WITHIN.
            harmonize: harmonize pixel values across different processing baselines,
                see https://developers.google.com/earth-engine/datasets/catalog/COPERNICUS_S2_SR_HARMONIZED
            context: the data source context.
        """  # noqa: E501
        # If context is provided, we join the directory with the dataset path,
        # otherwise we treat it directly as UPath.
        if context.ds_path is not None:
            self.metadata_cache_dir = join_upath(context.ds_path, metadata_cache_dir)
        else:
            self.metadata_cache_dir = UPath(metadata_cache_dir)

        self.modality = modality
        self.sort_by = sort_by
        self.harmonize = harmonize

        bucket_name = self.bucket_names[modality]
        self.bucket = boto3.resource(
            "s3", config=botocore.client.Config(signature_version=botocore.UNSIGNED)
        ).Bucket(bucket_name)

    def _read_products(
        self, needed_cell_months: set[tuple[str, int, int]]
    ) -> Generator[Sentinel2Item, None, None]:
        """Read productInfo.json files and yield relevant Sentinel2Items.

        Args:
            needed_cell_months: set of (mgrs grid cell, year, month) where we need
                to search for images.
        """
        for cell_id, year, month in tqdm.tqdm(
            needed_cell_months, desc="Reading product infos"
        ):
            assert len(cell_id) == 5
            cache_fname = self.metadata_cache_dir / f"{cell_id}_{year}_{month}.json"

            if not cache_fname.exists():
                cell_part1 = cell_id[0:2]
                cell_part2 = cell_id[2:3]
                cell_part3 = cell_id[3:5]
                prefix = (
                    f"tiles/{cell_part1}/{cell_part2}/{cell_part3}/"
                    + f"{year}/{month}/"
                )

                products = []
                for obj in self.bucket.objects.filter(Prefix=prefix):
                    if not obj.key.endswith("tileInfo.json"):
                        continue
                    buf = io.BytesIO()
                    self.bucket.download_fileobj(obj.key, buf)
                    buf.seek(0)
                    product = json.load(buf)
                    if "tileDataGeometry" not in product:
                        print(
                            "warning: skipping product missing tileDataGeometry",
                            product,
                        )
                        continue
                    if product["tileDataGeometry"]["type"] != "Polygon":
                        print(
                            "warning: skipping product with non-polygon geometry",
                            product,
                        )
                        continue
                    products.append(product)

                with open_atomic(cache_fname, "w") as f:
                    json.dump(products, f)

            else:
                with cache_fname.open() as f:
                    products = json.load(f)

            for product in products:
                tile_geometry_dict = product["tileDataGeometry"]
                assert tile_geometry_dict["crs"]["type"] == "name"
                crs = CRS.from_string(tile_geometry_dict["crs"]["properties"]["name"])
                ts = dateutil.parser.isoparse(product["timestamp"])
                geometry = STGeometry(
                    Projection(crs, 1, 1),
                    shapely.geometry.shape(tile_geometry_dict),
                    (ts, ts),
                )
                yield Sentinel2Item(
                    name=product["productName"],
                    geometry=geometry,
                    blob_path=product["path"] + "/",
                    cloud_cover=product["cloudyPixelPercentage"],
                )

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[Sentinel2Item]]]:
        """Get a list of items in the data source intersecting the given geometries.

        Args:
            geometries: the spatiotemporal geometries
            query_config: the query configuration

        Returns:
            List of groups of items that should be retrieved for each geometry.
        """
        # Identify all (mgrs grid cell, year, month, day) where we need to search for
        # images.
        # To do so, we iterate over the geometries and figure out all the MGRS cells
        # that they intersect.
        needed_cell_months = set()
        wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]
        for wgs84_geometry in wgs84_geometries:
            if wgs84_geometry.time_range is None:
                raise ValueError(
                    "Sentinel2 on AWS requires geometry time ranges to be set"
                )
            for cell_id in get_sentinel2_tiles(wgs84_geometry, self.metadata_cache_dir):
                for ts in daterange(
                    wgs84_geometry.time_range[0],
                    wgs84_geometry.time_range[1],
                ):
                    needed_cell_months.add((cell_id, ts.year, ts.month))

        items_by_cell: dict[str, list[Sentinel2Item]] = {}
        for item in self._read_products(needed_cell_months):
            cell_id = "".join(item.blob_path.split("/")[1:4])
            if cell_id not in items_by_cell:
                items_by_cell[cell_id] = []
            items_by_cell[cell_id].append(item)

        groups = []
        for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
            items = []
            for cell_id in get_sentinel2_tiles(wgs84_geometry, self.metadata_cache_dir):
                for item in items_by_cell.get(cell_id, []):
                    try:
                        item_geom = item.geometry.to_projection(geometry.projection)
                    except ValueError as e:
                        print(f"error re-projecting item {item.name}: {e}")
                        continue
                    if not geometry.shp.intersects(item_geom.shp):
                        continue
                    items.append(item)

            if self.sort_by == "cloud_cover":
                items.sort(key=lambda item: item.cloud_cover)
            elif self.sort_by is not None:
                raise ValueError(f"invalid sort_by setting ({self.sort_by})")

            cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
                geometry, items, query_config
            )
            groups.append(cur_groups)

        return groups

    def get_item_by_name(self, name: str) -> Sentinel2Item:
        """Gets an item by name."""
        # Product name is like:
        #   S2B_MSIL1C_20240201T230819_N0510_R015_T51CWM_20240202T012755.
        # We want to use _read_products so we need to extract:
        # - cell_id: 51CWM
        # - year: 2024
        # - month: 2
        parts = name.split("_")
        assert parts[5][0] == "T"
        cell_id = parts[5][1:]
        assert len(parts[2]) == 15
        year = int(parts[2][0:4])
        month = int(parts[2][4:6])
        for item in self._read_products({(cell_id, year, month)}):
            if item.name == name:
                return item
        raise ValueError(f"item {name} not found")

    def deserialize_item(self, serialized_item: dict) -> Sentinel2Item:
        """Deserializes an item from JSON-decoded data."""
        return Sentinel2Item.deserialize(serialized_item)

    def retrieve_item(
        self, item: Sentinel2Item
    ) -> Generator[tuple[str, BinaryIO], None, None]:
        """Retrieves the rasters corresponding to an item as file streams."""
        for fname, _ in self.band_fnames[self.modality]:
            buf = io.BytesIO()
            self.bucket.download_fileobj(item.blob_path + fname, buf)
            buf.seek(0)
            yield (fname, buf)

    def _get_harmonize_callback(
        self, item: Sentinel2Item
    ) -> Callable[[npt.NDArray], npt.NDArray] | None:
        """Gets the harmonization callback for the given item.

        Args:
            item: the item

        Returns:
            None if no callback is needed, or the callback to subtract the new offset
        """
        if not self.harmonize:
            return None
        # Search metadata XML for the RADIO_ADD_OFFSET tag.
        # This contains the per-band offset, but we assume all bands have the same offset.
        if item.geometry.time_range is None:
            raise ValueError("Sentinel2 on AWS requires geometry time ranges to be set")
        ts = item.geometry.time_range[0]
        metadata_fname = (
            f"products/{ts.year}/{ts.month}/{ts.day}/{item.name}/metadata.xml"
        )
        buf = io.BytesIO()
        self.bucket.download_fileobj(metadata_fname, buf)
        buf.seek(0)
        tree: ET.ElementTree[ET.Element[str]] = ET.ElementTree(
            ET.fromstring(buf.getvalue())
        )
        return get_harmonize_callback(tree)

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[Sentinel2Item],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        Args:
            tile_store: the tile store to ingest into
            items: the items to ingest
            geometries: a list of geometries needed for each item
        """
        for item in items:
            for fname, band_names in self.band_fnames[self.modality]:
                if tile_store.is_raster_ready(item, band_names):
                    continue

                with tempfile.TemporaryDirectory() as tmp_dir:
                    local_fname = os.path.join(tmp_dir, fname.split("/")[-1])

                    try:
                        self.bucket.download_file(
                            item.blob_path + fname,
                            local_fname,
                        )
                    except Exception as e:
                        # TODO: sometimes for some reason object doesn't exist
                        # we should probably investigate further why it happens
                        # and then should create the layer here and mark it completed
                        print(
                            f"warning: got error {e} downloading {item.blob_path + fname}"
                        )
                        continue

                    harmonize_callback = self._get_harmonize_callback(item)

                    if harmonize_callback is not None:
                        # In this case we need to read the array, convert the pixel
                        # values, and pass modified array directly to the TileStore.
                        with rasterio.open(local_fname) as src:
                            array = src.read()
                            src_nodata = src.nodata
                            projection, bounds = get_raster_projection_and_bounds(src)
                        array = harmonize_callback(array)
                        raster_metadata = RasterMetadata(nodata_value=src_nodata)
                        tile_store.write_raster(
                            item,
                            band_names,
                            projection,
                            bounds,
                            RasterArray(
                                chw_array=array,
                                time_range=item.geometry.time_range,
                                metadata=raster_metadata,
                            ),
                        )

                    else:
                        tile_store.write_raster_file(
                            item,
                            band_names,
                            UPath(local_fname),
                            time_range=item.geometry.time_range,
                        )

get_items

get_items(geometries: list[STGeometry], query_config: QueryConfig) -> list[list[MatchedItemGroup[Sentinel2Item]]]

Get a list of items in the data source intersecting the given geometries.

Parameters:

Name	Type	Description	Default
geometries	list[STGeometry]	the spatiotemporal geometries	required
query_config	QueryConfig	the query configuration	required

Returns:

Type	Description
list[list[MatchedItemGroup[Sentinel2Item]]]	List of groups of items that should be retrieved for each geometry.

Source code in rslearn/data_sources/aws_open_data.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[Sentinel2Item]]]:
    """Get a list of items in the data source intersecting the given geometries.

    Args:
        geometries: the spatiotemporal geometries
        query_config: the query configuration

    Returns:
        List of groups of items that should be retrieved for each geometry.
    """
    # Identify all (mgrs grid cell, year, month, day) where we need to search for
    # images.
    # To do so, we iterate over the geometries and figure out all the MGRS cells
    # that they intersect.
    needed_cell_months = set()
    wgs84_geometries = [geometry.to_wgs84() for geometry in geometries]
    for wgs84_geometry in wgs84_geometries:
        if wgs84_geometry.time_range is None:
            raise ValueError(
                "Sentinel2 on AWS requires geometry time ranges to be set"
            )
        for cell_id in get_sentinel2_tiles(wgs84_geometry, self.metadata_cache_dir):
            for ts in daterange(
                wgs84_geometry.time_range[0],
                wgs84_geometry.time_range[1],
            ):
                needed_cell_months.add((cell_id, ts.year, ts.month))

    items_by_cell: dict[str, list[Sentinel2Item]] = {}
    for item in self._read_products(needed_cell_months):
        cell_id = "".join(item.blob_path.split("/")[1:4])
        if cell_id not in items_by_cell:
            items_by_cell[cell_id] = []
        items_by_cell[cell_id].append(item)

    groups = []
    for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
        items = []
        for cell_id in get_sentinel2_tiles(wgs84_geometry, self.metadata_cache_dir):
            for item in items_by_cell.get(cell_id, []):
                try:
                    item_geom = item.geometry.to_projection(geometry.projection)
                except ValueError as e:
                    print(f"error re-projecting item {item.name}: {e}")
                    continue
                if not geometry.shp.intersects(item_geom.shp):
                    continue
                items.append(item)

        if self.sort_by == "cloud_cover":
            items.sort(key=lambda item: item.cloud_cover)
        elif self.sort_by is not None:
            raise ValueError(f"invalid sort_by setting ({self.sort_by})")

        cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
            geometry, items, query_config
        )
        groups.append(cur_groups)

    return groups

get_item_by_name

get_item_by_name(name: str) -> Sentinel2Item

Gets an item by name.

Source code in rslearn/data_sources/aws_open_data.py

def get_item_by_name(self, name: str) -> Sentinel2Item:
    """Gets an item by name."""
    # Product name is like:
    #   S2B_MSIL1C_20240201T230819_N0510_R015_T51CWM_20240202T012755.
    # We want to use _read_products so we need to extract:
    # - cell_id: 51CWM
    # - year: 2024
    # - month: 2
    parts = name.split("_")
    assert parts[5][0] == "T"
    cell_id = parts[5][1:]
    assert len(parts[2]) == 15
    year = int(parts[2][0:4])
    month = int(parts[2][4:6])
    for item in self._read_products({(cell_id, year, month)}):
        if item.name == name:
            return item
    raise ValueError(f"item {name} not found")

deserialize_item

deserialize_item(serialized_item: dict) -> Sentinel2Item

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/aws_open_data.py

def deserialize_item(self, serialized_item: dict) -> Sentinel2Item:
    """Deserializes an item from JSON-decoded data."""
    return Sentinel2Item.deserialize(serialized_item)

retrieve_item

retrieve_item(item: Sentinel2Item) -> Generator[tuple[str, BinaryIO], None, None]

Retrieves the rasters corresponding to an item as file streams.

Source code in rslearn/data_sources/aws_open_data.py

def retrieve_item(
    self, item: Sentinel2Item
) -> Generator[tuple[str, BinaryIO], None, None]:
    """Retrieves the rasters corresponding to an item as file streams."""
    for fname, _ in self.band_fnames[self.modality]:
        buf = io.BytesIO()
        self.bucket.download_fileobj(item.blob_path + fname, buf)
        buf.seek(0)
        yield (fname, buf)

ingest

ingest(tile_store: TileStoreWithLayer, items: list[Sentinel2Item], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Parameters:

Name	Type	Description	Default
tile_store	TileStoreWithLayer	the tile store to ingest into	required
items	list[Sentinel2Item]	the items to ingest	required
geometries	list[list[STGeometry]]	a list of geometries needed for each item	required

Source code in rslearn/data_sources/aws_open_data.py

def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[Sentinel2Item],
    geometries: list[list[STGeometry]],
) -> None:
    """Ingest items into the given tile store.

    Args:
        tile_store: the tile store to ingest into
        items: the items to ingest
        geometries: a list of geometries needed for each item
    """
    for item in items:
        for fname, band_names in self.band_fnames[self.modality]:
            if tile_store.is_raster_ready(item, band_names):
                continue

            with tempfile.TemporaryDirectory() as tmp_dir:
                local_fname = os.path.join(tmp_dir, fname.split("/")[-1])

                try:
                    self.bucket.download_file(
                        item.blob_path + fname,
                        local_fname,
                    )
                except Exception as e:
                    # TODO: sometimes for some reason object doesn't exist
                    # we should probably investigate further why it happens
                    # and then should create the layer here and mark it completed
                    print(
                        f"warning: got error {e} downloading {item.blob_path + fname}"
                    )
                    continue

                harmonize_callback = self._get_harmonize_callback(item)

                if harmonize_callback is not None:
                    # In this case we need to read the array, convert the pixel
                    # values, and pass modified array directly to the TileStore.
                    with rasterio.open(local_fname) as src:
                        array = src.read()
                        src_nodata = src.nodata
                        projection, bounds = get_raster_projection_and_bounds(src)
                    array = harmonize_callback(array)
                    raster_metadata = RasterMetadata(nodata_value=src_nodata)
                    tile_store.write_raster(
                        item,
                        band_names,
                        projection,
                        bounds,
                        RasterArray(
                            chw_array=array,
                            time_range=item.geometry.time_range,
                            metadata=raster_metadata,
                        ),
                    )

                else:
                    tile_store.write_raster_file(
                        item,
                        band_names,
                        UPath(local_fname),
                        time_range=item.geometry.time_range,
                    )