rslearn 0.0.1__tar.gz → 0.0.2__tar.gz

{rslearn-0.0.1/rslearn.egg-info → rslearn-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rslearn
-Version: 0.0.1
+Version: 0.0.2
 Summary: A library for developing remote sensing datasets and models
 Author-email: Favyen Bastani <favyenb@allenai.org>, Yawen Zhang <yawenz@allenai.org>, Patrick Beukema <patrickb@allenai.org>, Henry Herzog <henryh@allenai.org>, Piper Wolters <piperw@allenai.org>
 License:                                  Apache License
@@ -208,33 +208,38 @@ License:                                  Apache License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: boto3
-Requires-Dist: class_registry
-Requires-Dist: python-dateutil
-Requires-Dist: pytimeparse
-Requires-Dist: fiona
-Requires-Dist: fsspec[gcs,s3]
-Requires-Dist: Pillow
-Requires-Dist: pyproj
-Requires-Dist: rasterio
-Requires-Dist: shapely
-Requires-Dist: tqdm
-Requires-Dist: torch
-Requires-Dist: torchvision
-Requires-Dist: universal_pathlib
-Requires-Dist: lightning[pytorch-extra]
+Requires-Dist: boto3>=1.35
+Requires-Dist: class_registry>=2.1
+Requires-Dist: fiona>=1.10
+Requires-Dist: fsspec[gcs,s3]>=2024.10
+Requires-Dist: lightning[pytorch-extra]>=2.4
+Requires-Dist: Pillow>=11.0
+Requires-Dist: pyproj>=3.7
+Requires-Dist: python-dateutil>=2.9
+Requires-Dist: pytimeparse>=1.1
+Requires-Dist: rasterio>=1.4
+Requires-Dist: shapely>=2.0
+Requires-Dist: torch>=2.5
+Requires-Dist: torchvision>=0.20
+Requires-Dist: tqdm>=4.66
+Requires-Dist: universal_pathlib>=0.2.5
 Provides-Extra: extra
-Requires-Dist: earthengine-api; extra == "extra"
-Requires-Dist: gcsfs; extra == "extra"
-Requires-Dist: google-cloud-storage; extra == "extra"
-Requires-Dist: mgrs; extra == "extra"
-Requires-Dist: osmium; extra == "extra"
-Requires-Dist: planet; extra == "extra"
-Requires-Dist: pycocotools; extra == "extra"
-Requires-Dist: rtree; extra == "extra"
-Requires-Dist: satlaspretrain_models; extra == "extra"
-Requires-Dist: scipy; extra == "extra"
-Requires-Dist: wandb; extra == "extra"
+Requires-Dist: accelerate>=1.0; extra == "extra"
+Requires-Dist: earthengine-api>=0.1; extra == "extra"
+Requires-Dist: einops>=0.8; extra == "extra"
+Requires-Dist: gcsfs>=2024.10; extra == "extra"
+Requires-Dist: google-cloud-bigquery>=2.18; extra == "extra"
+Requires-Dist: google-cloud-storage>=2.18; extra == "extra"
+Requires-Dist: interrogate>=1.7; extra == "extra"
+Requires-Dist: osmium>=3.7; extra == "extra"
+Requires-Dist: planet>=2.10; extra == "extra"
+Requires-Dist: pycocotools>=2.0; extra == "extra"
+Requires-Dist: rtree>=1.2; extra == "extra"
+Requires-Dist: s3fs>=2024.10.0; extra == "extra"
+Requires-Dist: satlaspretrain_models>=0.3; extra == "extra"
+Requires-Dist: scipy>=1.13; extra == "extra"
+Requires-Dist: transformers>=4.45; extra == "extra"
+Requires-Dist: wandb>=0.17; extra == "extra"
 
 Overview
 --------

rslearn-0.0.2/extra_requirements.txt

@@ -0,0 +1,17 @@
+# These are requirements that are specific to a subset of data sources or models.
+accelerate>=1.0
+earthengine-api>=0.1
+einops>=0.8
+gcsfs>=2024.10
+google-cloud-bigquery>=2.18
+google-cloud-storage>=2.18
+interrogate>=1.7
+osmium>=3.7
+planet>=2.10
+pycocotools>=2.0
+rtree>=1.2
+s3fs>=2024.10.0
+satlaspretrain_models>=0.3
+scipy>=1.13
+transformers>=4.45
+wandb>=0.17

{rslearn-0.0.1 → rslearn-0.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.1"
+version = "0.0.2"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     {name = "Favyen Bastani", email = "favyenb@allenai.org"},

rslearn-0.0.2/requirements.txt

@@ -0,0 +1,15 @@
+boto3>=1.35
+class_registry>=2.1
+fiona>=1.10
+fsspec[s3,gcs]>=2024.10
+lightning[pytorch-extra]>=2.4
+Pillow>=11.0
+pyproj>=3.7
+python-dateutil>=2.9
+pytimeparse>=1.1
+rasterio>=1.4
+shapely>=2.0
+torch>=2.5
+torchvision>=0.20
+tqdm>=4.66
+universal_pathlib>=0.2.5

{rslearn-0.0.1 → rslearn-0.0.2}/rslearn/config/dataset.py

@@ -112,7 +112,7 @@ class BandSetConfig:
         self,
         config_dict: dict[str, Any],
         dtype: DType,
-        bands: list[str] | None = None,
+        bands: list[str],
         format: dict[str, Any] | None = None,
         zoom_offset: int = 0,
         remap: dict[str, Any] | None = None,
@@ -130,13 +130,14 @@ class BandSetConfig:
         """
         self.config_dict = config_dict
         self.bands = bands
-        self.format = format
         self.dtype = dtype
         self.zoom_offset = zoom_offset
         self.remap = remap
 
-        if not self.format:
+        if format is None:
             self.format = {"name": "geotiff"}
+        else:
+            self.format = format
 
     def serialize(self) -> dict[str, Any]:
         """Serialize this BandSetConfig to a config dict, currently unused."""
@@ -158,11 +159,12 @@ class BandSetConfig:
         kwargs = dict(
             config_dict=config,
             dtype=DType(config["dtype"]),
+            bands=config["bands"],
         )
-        for k in ["bands", "format", "zoom_offset", "remap"]:
+        for k in ["format", "zoom_offset", "remap"]:
             if k in config:
                 kwargs[k] = config[k]
-        return BandSetConfig(**kwargs)
+        return BandSetConfig(**kwargs)  # type: ignore
 
     def get_final_projection_and_bounds(
         self, projection: Projection, bounds: PixelBounds | None
@@ -187,11 +189,15 @@ class BandSetConfig:
             projection.x_resolution / (2**self.zoom_offset),
             projection.y_resolution / (2**self.zoom_offset),
         )
-        if bounds:
+        if bounds is not None:
             if self.zoom_offset > 0:
-                bounds = tuple(x * (2**self.zoom_offset) for x in bounds)
+                zoom_factor = 2**self.zoom_offset
+                bounds = tuple(x * zoom_factor for x in bounds)  # type: ignore
             else:
-                bounds = tuple(x // (2 ** (-self.zoom_offset)) for x in bounds)
+                bounds = tuple(
+                    x // (2 ** (-self.zoom_offset))
+                    for x in bounds  # type: ignore
+                )
         return projection, bounds
 
 
@@ -422,7 +428,7 @@ class RasterLayerConfig(LayerConfig):
             ]
         if "alias" in config:
             kwargs["alias"] = config["alias"]
-        return RasterLayerConfig(**kwargs)
+        return RasterLayerConfig(**kwargs)  # type: ignore
 
 
 class VectorLayerConfig(LayerConfig):
@@ -456,7 +462,7 @@ class VectorLayerConfig(LayerConfig):
         Args:
             config: the config dict for this VectorLayerConfig
         """
-        kwargs = {"layer_type": LayerType(config["type"])}
+        kwargs: dict[str, Any] = {"layer_type": LayerType(config["type"])}
         if "data_source" in config:
             kwargs["data_source"] = DataSourceConfig.from_config(config["data_source"])
         if "zoom_offset" in config:
@@ -465,7 +471,7 @@ class VectorLayerConfig(LayerConfig):
             kwargs["format"] = VectorFormatConfig.from_config(config["format"])
         if "alias" in config:
             kwargs["alias"] = config["alias"]
-        return VectorLayerConfig(**kwargs)
+        return VectorLayerConfig(**kwargs)  # type: ignore
 
     def get_final_projection_and_bounds(
         self, projection: Projection, bounds: PixelBounds | None
@@ -488,9 +494,12 @@ class VectorLayerConfig(LayerConfig):
         )
         if bounds:
             if self.zoom_offset > 0:
-                bounds = tuple(x * (2**self.zoom_offset) for x in bounds)
+                bounds = tuple(x * (2**self.zoom_offset) for x in bounds)  # type: ignore
             else:
-                bounds = tuple(x // (2 ** (-self.zoom_offset)) for x in bounds)
+                bounds = tuple(
+                    x // (2 ** (-self.zoom_offset))
+                    for x in bounds  # type: ignore
+                )
         return projection, bounds
 
 

{rslearn-0.0.1 → rslearn-0.0.2}/rslearn/data_sources/__init__.py

@@ -10,15 +10,20 @@ Each source supports operations to lookup items that match with spatiotemporal
 geometries, and ingest those items.
 """
 
+import functools
 import importlib
 
 from upath import UPath
 
 from rslearn.config import LayerConfig
+from rslearn.log_utils import get_logger
 
 from .data_source import DataSource, Item, ItemLookupDataSource, RetrieveItemDataSource
 
+logger = get_logger(__name__)
 
+
+@functools.cache
 def data_source_from_config(config: LayerConfig, ds_path: UPath) -> DataSource:
     """Loads a data source from config dict.
 
@@ -26,6 +31,9 @@ def data_source_from_config(config: LayerConfig, ds_path: UPath) -> DataSource:
         config: the LayerConfig containing this data source.
         ds_path: the dataset root directory.
     """
+    logger.debug("getting a data source for dataset at %s", ds_path)
+    if config.data_source is None:
+        raise ValueError("No data source specified")
     name = config.data_source.name
     module_name = ".".join(name.split(".")[:-1])
     class_name = name.split(".")[-1]

{rslearn-0.0.1 → rslearn-0.0.2}/rslearn/data_sources/aws_landsat.py

@@ -6,7 +6,7 @@ import shutil
 import urllib.request
 import zipfile
 from collections.abc import Generator
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Any, BinaryIO
 
 import boto3
@@ -20,8 +20,7 @@ import tqdm
 from upath import UPath
 
 import rslearn.data_sources.utils
-import rslearn.utils.mgrs
-from rslearn.config import LayerConfig, RasterLayerConfig
+from rslearn.config import RasterLayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_PROJECTION
 from rslearn.tile_stores import PrefixedTileStore, TileStore
 from rslearn.utils import STGeometry
@@ -36,7 +35,7 @@ class LandsatOliTirsItem(Item):
 
     def __init__(
         self, name: str, geometry: STGeometry, blob_path: str, cloud_cover: float
-    ):
+    ) -> None:
         """Creates a new LandsatOliTirsItem.
 
         Args:
@@ -58,7 +57,7 @@ class LandsatOliTirsItem(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "LandsatOliTirsItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         if "name" not in d:
             d["name"] = d["blob_path"].split("/")[-1].split(".tif")[0]
@@ -90,7 +89,7 @@ class LandsatOliTirs(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         metadata_cache_dir: UPath,
         max_time_delta: timedelta = timedelta(days=30),
         sort_by: str | None = None,
@@ -110,15 +109,16 @@ class LandsatOliTirs(DataSource):
         self.metadata_cache_dir = metadata_cache_dir
         self.max_time_delta = max_time_delta
         self.sort_by = sort_by
+        print(self.metadata_cache_dir)
 
         self.bucket = boto3.resource("s3").Bucket(self.bucket_name)
-
         self.metadata_cache_dir.mkdir(parents=True, exist_ok=True)
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LandsatOliTirs":
+    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "LandsatOliTirs":
         """Creates a new LandsatOliTirs instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError(f"data_source is required for config dict {config}")
         d = config.data_source.config_dict
         kwargs = dict(
             config=config,
@@ -181,8 +181,9 @@ class LandsatOliTirs(DataSource):
                     ts = dateutil.parser.isoparse(date_str + "T" + time_str)
 
                     blob_path = obj.key.split("MTL.json")[0]
+                    time_range: tuple[datetime, datetime] = (ts, ts)
                     geometry = STGeometry(
-                        WGS84_PROJECTION, shapely.Polygon(coordinates), [ts, ts]
+                        WGS84_PROJECTION, shapely.Polygon(coordinates), time_range
                     )
                     items.append(
                         LandsatOliTirsItem(
@@ -216,6 +217,7 @@ class LandsatOliTirs(DataSource):
         if not shp_fname.exists():
             # Download and extract zip to cache dir.
             zip_fname = self.metadata_cache_dir / f"{prefix}.zip"
+            print(f"Downloading {self.wrs2_url} to {zip_fname}")
             with urllib.request.urlopen(self.wrs2_url) as response:
                 with zip_fname.open("wb") as f:
                     shutil.copyfileobj(response, f)
@@ -259,7 +261,7 @@ class LandsatOliTirs(DataSource):
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[LandsatOliTirsItem]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         Args:
@@ -305,14 +307,16 @@ class LandsatOliTirs(DataSource):
             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, cur_items, query_config
+            cur_groups: list[list[LandsatOliTirsItem]] = (
+                rslearn.data_sources.utils.match_candidate_items_to_window(
+                    geometry, cur_items, query_config
+                )
             )
             groups.append(cur_groups)
 
         return groups
 
-    def get_item_by_name(self, name: str) -> Item:
+    def get_item_by_name(self, name: str) -> LandsatOliTirsItem:
         """Gets an item by name."""
         # Product name is like LC08_L1TP_046027_20230715_20230724_02_T1.
         # We want to use _read_products so we need to extract:
@@ -330,12 +334,14 @@ class LandsatOliTirs(DataSource):
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> LandsatOliTirsItem:
         """Deserializes an item from JSON-decoded data."""
         assert isinstance(serialized_item, dict)
         return LandsatOliTirsItem.deserialize(serialized_item)
 
-    def retrieve_item(self, item: Item) -> Generator[tuple[str, BinaryIO], None, None]:
+    def retrieve_item(
+        self, item: LandsatOliTirsItem
+    ) -> Generator[tuple[str, BinaryIO], None, None]:
         """Retrieves the rasters corresponding to an item as file streams."""
         for band in self.bands:
             buf = io.BytesIO()
@@ -351,7 +357,7 @@ class LandsatOliTirs(DataSource):
     def ingest(
         self,
         tile_store: TileStore,
-        items: list[Item],
+        items: list[LandsatOliTirsItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -368,7 +374,10 @@ class LandsatOliTirs(DataSource):
                     tile_store, (item.name, "_".join(band_names))
                 )
                 needed_projections = get_needed_projections(
-                    cur_tile_store, band_names, self.config.band_sets, cur_geometries
+                    cur_tile_store,
+                    band_names,
+                    self.config.band_sets,
+                    cur_geometries,  # type: ignore
                 )
                 if not needed_projections:
                     continue

{rslearn-0.0.1 → rslearn-0.0.2}/rslearn/data_sources/aws_open_data.py

@@ -2,7 +2,6 @@
 
 import io
 import json
-import tempfile
 import xml.etree.ElementTree as ET
 from collections.abc import Callable, Generator
 from datetime import datetime, timedelta, timezone
@@ -22,19 +21,13 @@ from rasterio.crs import CRS
 from upath import UPath
 
 import rslearn.data_sources.utils
-import rslearn.utils.mgrs
-from rslearn.config import LayerConfig, RasterLayerConfig
+from rslearn.config import RasterLayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_EPSG, WGS84_PROJECTION
 from rslearn.tile_stores import PrefixedTileStore, TileStore
-from rslearn.utils import (
-    GridIndex,
-    Projection,
-    STGeometry,
-    daterange,
-)
+from rslearn.utils import GridIndex, Projection, STGeometry, daterange
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
 
-from .copernicus import get_harmonize_callback
+from .copernicus import get_harmonize_callback, get_sentinel2_tiles
 from .data_source import (
     DataSource,
     Item,
@@ -66,7 +59,7 @@ class NaipItem(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    def deserialize(d: dict) -> "NaipItem":
         """Deserializes an item from a JSON-decoded dictionary."""
         item = super(NaipItem, NaipItem).deserialize(d)
         return NaipItem(
@@ -89,7 +82,7 @@ class Naip(DataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         index_cache_dir: UPath,
         use_rtree_index: bool = False,
         states: list[str] | None = None,
@@ -113,25 +106,21 @@ class Naip(DataSource):
         self.years = years
 
         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):
+            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_tmp_dir = tempfile.TemporaryDirectory()
-            self.rtree_index = get_cached_rtree(
-                self.index_cache_dir, self.rtree_tmp_dir.name, build_fn
-            )
-        else:
-            self.rtree_index = None
+            self.rtree_index = get_cached_rtree(self.index_cache_dir, build_fn)
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "Naip":
+    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Naip":
         """Creates a new Naip instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError(f"data_source is required for config dict {config}")
         d = config.data_source.config_dict
         kwargs = dict(
             config=config,
@@ -195,7 +184,9 @@ class Naip(DataSource):
                     blob_path, dst, ExtraArgs={"RequestPayer": "requester"}
                 )
 
-    def _read_index_shapefiles(self, desc=None) -> Generator[NaipItem, None, None]:
+    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()
 
@@ -288,7 +279,7 @@ class Naip(DataSource):
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[NaipItem]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         Args:
@@ -302,7 +293,7 @@ class Naip(DataSource):
             geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
         ]
 
-        items = [[] for _ in geometries]
+        items: list = [[] for _ in geometries]
         if self.rtree_index:
             for idx, geometry in enumerate(wgs84_geometries):
                 encoded_items = self.rtree_index.query(geometry.shp.bounds)
@@ -331,7 +322,7 @@ class Naip(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> NaipItem:
         """Deserializes an item from JSON-decoded data."""
         assert isinstance(serialized_item, dict)
         return NaipItem.deserialize(serialized_item)
@@ -339,7 +330,7 @@ class Naip(DataSource):
     def ingest(
         self,
         tile_store: TileStore,
-        items: list[Item],
+        items: list[NaipItem],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.
@@ -407,7 +398,7 @@ class Sentinel2Item(Item):
         return d
 
     @staticmethod
-    def deserialize(d: dict) -> Item:
+    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]
@@ -420,7 +411,10 @@ class Sentinel2Item(Item):
         )
 
 
-class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
+# TODO: Distinguish between AWS and GCP data sources in class names.
+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
@@ -474,7 +468,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
 
     def __init__(
         self,
-        config: LayerConfig,
+        config: RasterLayerConfig,
         modality: Sentinel2Modality,
         metadata_cache_dir: UPath,
         max_time_delta: timedelta = timedelta(days=30),
@@ -506,9 +500,10 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
         self.bucket = boto3.resource("s3").Bucket(bucket_name)
 
     @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "Sentinel2":
+    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Sentinel2":
         """Creates a new Sentinel2 instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
+        if config.data_source is None:
+            raise ValueError("Sentinel2 data source requires a data source config")
         d = config.data_source.config_dict
         kwargs = dict(
             config=config,
@@ -528,7 +523,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
         return Sentinel2(**kwargs)
 
     def _read_products(
-        self, needed_cell_months: set[tuple[str, int, int, int]]
+        self, needed_cell_months: set[tuple[str, int, int]]
     ) -> Generator[Sentinel2Item, None, None]:
         """Read productInfo.json files and yield relevant Sentinel2Items.
 
@@ -603,7 +598,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
 
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
-    ) -> list[list[list[Item]]]:
+    ) -> list[list[list[Sentinel2Item]]]:
         """Get a list of items in the data source intersecting the given geometries.
 
         Args:
@@ -626,14 +621,14 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
                 raise ValueError(
                     "Sentinel2 on AWS requires geometry time ranges to be set"
                 )
-            for cell_id in rslearn.utils.mgrs.for_each_cell(wgs84_geometry.shp.bounds):
+            for cell_id in get_sentinel2_tiles(wgs84_geometry, self.metadata_cache_dir):
                 for ts in daterange(
                     wgs84_geometry.time_range[0] - self.max_time_delta,
                     wgs84_geometry.time_range[1] + self.max_time_delta,
                 ):
                     needed_cell_months.add((cell_id, ts.year, ts.month))
 
-        items_by_cell = {}
+        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:
@@ -643,7 +638,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
         groups = []
         for geometry, wgs84_geometry in zip(geometries, wgs84_geometries):
             items = []
-            for cell_id in rslearn.utils.mgrs.for_each_cell(wgs84_geometry.shp.bounds):
+            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)
@@ -666,7 +661,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
 
         return groups
 
-    def get_item_by_name(self, name: str) -> Item:
+    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.
@@ -685,12 +680,14 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: Any) -> Sentinel2Item:
         """Deserializes an item from JSON-decoded data."""
         assert isinstance(serialized_item, dict)
         return Sentinel2Item.deserialize(serialized_item)
 
-    def retrieve_item(self, item: Item) -> Generator[tuple[str, BinaryIO], None, None]:
+    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()
@@ -701,7 +698,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
             yield (fname, buf)
 
     def _get_harmonize_callback(
-        self, item: Item
+        self, item: Sentinel2Item
     ) -> Callable[[npt.NDArray], npt.NDArray] | None:
         """Gets the harmonization callback for the given item.
 
@@ -715,6 +712,8 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
             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"
@@ -730,7 +729,7 @@ class Sentinel2(ItemLookupDataSource, RetrieveItemDataSource):
     def ingest(
         self,
         tile_store: TileStore,
-        items: list[Item],
+        items: list[Sentinel2Item],
         geometries: list[list[STGeometry]],
     ) -> None:
         """Ingest items into the given tile store.