rslearn 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

rslearn/data_sources/earthdaily.py

@@ -17,9 +17,9 @@ from earthdaily import EDSClient, EDSConfig
 from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import LayerConfig, QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.dataset import Window
 from rslearn.dataset.materialize import RasterMaterializer
@@ -81,10 +81,11 @@ class EarthDaily(DataSource, TileStore):
         sort_ascending: bool = True,
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
-        cache_dir: UPath | None = None,
+        cache_dir: str | None = None,
         max_retries: int = 3,
         retry_backoff_factor: float = 5.0,
         service_name: Literal["platform"] = "platform",
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new EarthDaily instance.
 
@@ -108,6 +109,7 @@ class EarthDaily(DataSource, TileStore):
                 `(retry_backoff_factor * (2 ** (retry_count - 1)))` seconds.
             service_name: the service name, only "platform" is supported, the other
                 services "legacy" and "internal" are not supported.
+            context: the data source context.
         """
         self.collection_name = collection_name
         self.asset_bands = asset_bands
@@ -116,51 +118,25 @@ class EarthDaily(DataSource, TileStore):
         self.sort_ascending = sort_ascending
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
-        self.cache_dir = cache_dir
         self.max_retries = max_retries
         self.retry_backoff_factor = retry_backoff_factor
         self.service_name = service_name
 
         if cache_dir is not None:
-            self.cache_dir = cache_dir
+            # Use dataset path as root if provided.
+            if context.ds_path is not None:
+                self.cache_dir = join_upath(context.ds_path, cache_dir)
+            else:
+                self.cache_dir = UPath(cache_dir)
+
             self.cache_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            self.cache_dir = None
 
         self.eds_client: EDSClient | None = None
         self.client: pystac_client.Client | None = None
         self.collection: pystac_client.CollectionClient | None = None
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "EarthDaily":
-        """Creates a new EarthDaily instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-
-        kwargs: dict[str, Any] = dict(
-            collection_name=d["collection_name"],
-            service_name=d["service_name"],
-            asset_bands=d["asset_bands"],
-        )
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-
-        if "max_retries" in d:
-            kwargs["max_retries"] = d["max_retries"]
-
-        if "retry_backoff_factor" in d:
-            kwargs["retry_backoff_factor"] = d["retry_backoff_factor"]
-
-        simple_optionals = ["query", "sort_by", "sort_ascending"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return EarthDaily(**kwargs)
-
     def _load_client(
         self,
     ) -> tuple[EDSClient, pystac_client.Client, pystac_client.CollectionClient]:
@@ -499,7 +475,6 @@ class EarthDaily(DataSource, TileStore):
             layer_name: the name of this layer
             layer_cfg: the config of this layer
         """
-        assert isinstance(layer_cfg, RasterLayerConfig)
         RasterMaterializer().materialize(
             TileStoreWithLayer(self, layer_name),
             window,

rslearn/data_sources/earthdata_srtm.py

@@ -12,9 +12,9 @@ import requests.auth
 import shapely
 from upath import UPath
 
-from rslearn.config import QueryConfig, RasterLayerConfig, SpaceMode
+from rslearn.config import QueryConfig, SpaceMode
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.geometry import STGeometry
@@ -47,8 +47,8 @@ class SRTM(DataSource):
         self,
         username: str | None = None,
         password: str | None = None,
-        band_name: str = "srtm",
         timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new SRTM instance.
 
@@ -57,10 +57,19 @@ class SRTM(DataSource):
                 NASA_EARTHDATA_USERNAME environment variable.
             password: NASA Earthdata account password. If not set, it is read from the
                 NASA_EARTHDATA_PASSWORD environment variable.
-            band_name: what to call the band.
             timeout: timeout for requests.
+            context: the data source context.
         """
-        self.band_name = band_name
+        # Get band name from context if possible, falling back to "srtm".
+        if context.layer_config is not None:
+            if len(context.layer_config.band_sets) != 1:
+                raise ValueError("expected a single band set")
+            if len(context.layer_config.band_sets[0].bands) != 1:
+                raise ValueError("expected band set to have a single band")
+            self.band_name = context.layer_config.band_sets[0].bands[0]
+        else:
+            self.band_name = "srtm"
+
         self.timeout = timeout
 
         if username is None:
@@ -73,33 +82,6 @@ class SRTM(DataSource):
 
         self.session = requests.session()
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "SRTM":
-        """Creates a new SRTM instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-
-        # Get the band name chosen by the user.
-        # There should be a single band set with a single band.
-        if len(config.band_sets) != 1:
-            raise ValueError("expected a single band set")
-        if len(config.band_sets[0].bands) != 1:
-            raise ValueError("expected band set to have a single band")
-        kwargs: dict[str, Any] = {
-            "band_name": config.band_sets[0].bands[0],
-        }
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        simple_optionals = ["username", "password"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return SRTM(**kwargs)
-
     def get_item_by_name(self, name: str) -> Item:
         """Gets an item by name.
 

rslearn/data_sources/eurocrops.py

@@ -10,11 +10,10 @@ from typing import Any
 import fiona
 import requests
 from rasterio.crs import CRS
-from upath import UPath
 
-from rslearn.config import QueryConfig, VectorLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
@@ -108,12 +107,9 @@ class EuroCrops(DataSource[EuroCropsItem]):
     }
     TIMEOUT = timedelta(seconds=10)
 
-    @staticmethod
-    def from_config(config: VectorLayerConfig, ds_path: UPath) -> "EuroCrops":
-        """Creates a new EuroCrops instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("data_source is required")
-        return EuroCrops()
+    def __init__(self, context: DataSourceContext = DataSourceContext()):
+        """Create a new EuroCrops."""
+        pass
 
     def _get_all_items(self) -> list[EuroCropsItem]:
         """Get a list of all available items in the data source."""

rslearn/data_sources/gcp_public_data.py

@@ -18,10 +18,9 @@ import tqdm
 from google.cloud import bigquery, storage
 from upath import UPath
 
-from rslearn.config import QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
-from rslearn.data_sources.raster_source import is_raster_needed
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
@@ -157,19 +156,19 @@ class Sentinel2(DataSource):
 
     def __init__(
         self,
-        config: RasterLayerConfig,
-        index_cache_dir: UPath,
+        index_cache_dir: str,
         sort_by: str | None = None,
         use_rtree_index: bool = True,
         harmonize: bool = False,
         rtree_time_range: tuple[datetime, datetime] | None = None,
-        rtree_cache_dir: UPath | None = None,
+        rtree_cache_dir: str | None = None,
         use_bigquery: bool | None = None,
+        bands: list[str] | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new Sentinel2 instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source.
             index_cache_dir: local directory to cache the index contents, as well as
                 individual product metadata files.
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
@@ -193,6 +192,9 @@ class Sentinel2(DataSource):
                 credentials, set use_bigquery=False and use_rtree_index=False. The
                 default value is None which enables BigQuery when use_rtree_index=True
                 and disables when use_rtree_index=False.
+            bands: the bands to download, or None to download all bands. This is only
+                used if the layer config is not in the context.
+            context: the data source context.
         """
         if use_bigquery is None:
             use_bigquery = use_rtree_index
@@ -201,22 +203,52 @@ class Sentinel2(DataSource):
                 "use_bigquery must be enabled if use_rtree_index is enabled"
             )
 
-        self.config = config
-        self.index_cache_dir = index_cache_dir
+        # Resolve index_cache_dir and rtree_cache_dir depending on dataset context.
+        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)
+
+        if rtree_cache_dir is None:
+            self.rtree_cache_dir = self.index_cache_dir
+        elif context.ds_path is not None:
+            self.rtree_cache_dir = join_upath(context.ds_path, rtree_cache_dir)
+        else:
+            self.rtree_cache_dir = UPath(rtree_cache_dir)
+
         self.sort_by = sort_by
         self.harmonize = harmonize
         self.use_bigquery = use_bigquery
 
         self.index_cache_dir.mkdir(parents=True, exist_ok=True)
 
+        # Determine the subset of bands that are needed based on the layer config.
+        self.needed_bands: list[tuple[str, list[str]]]
+        if context.layer_config is not None:
+            self.needed_bands = []
+            for fname, cur_bands in self.BANDS:
+                # See if the bands provided by this file intersect with the bands in at
+                # least one configured band set.
+                for band_set in context.layer_config.band_sets:
+                    if not set(band_set.bands).intersection(cur_bands):
+                        continue
+                    self.needed_bands.append((fname, cur_bands))
+                    break
+        elif bands is not None:
+            self.needed_bands = []
+            for fname, cur_bands in self.BANDS:
+                if not set(bands).intersection(cur_bands):
+                    continue
+                self.needed_bands.append((fname, cur_bands))
+        else:
+            self.needed_bands = list(self.BANDS)
+
         self.bucket = storage.Client.create_anonymous_client().bucket(self.BUCKET_NAME)
         self.rtree_index: Any | None = None
         if use_rtree_index:
             from rslearn.utils.rtree_index import RtreeIndex, get_cached_rtree
 
-            if rtree_cache_dir is None:
-                rtree_cache_dir = self.index_cache_dir
-            rtree_cache_dir.mkdir(parents=True, exist_ok=True)
+            self.rtree_cache_dir.mkdir(parents=True, exist_ok=True)
 
             def build_fn(index: RtreeIndex) -> None:
                 """Build the RtreeIndex from items in the data source."""
@@ -226,34 +258,7 @@ class Sentinel2(DataSource):
                     for shp in flatten_shape(item.geometry.shp):
                         index.insert(shp.bounds, json.dumps(item.serialize()))
 
-            self.rtree_index = get_cached_rtree(rtree_cache_dir, build_fn)
-
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Sentinel2":
-        """Creates a new Sentinel2 instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            index_cache_dir=join_upath(ds_path, d["index_cache_dir"]),
-        )
-
-        if "rtree_time_range" in d:
-            kwargs["rtree_time_range"] = (
-                datetime.fromisoformat(d["rtree_time_range"][0]),
-                datetime.fromisoformat(d["rtree_time_range"][1]),
-            )
-
-        if "rtree_cache_dir" in d:
-            kwargs["rtree_cache_dir"] = join_upath(ds_path, d["rtree_cache_dir"])
-
-        simple_optionals = ["sort_by", "use_rtree_index", "harmonize", "use_bigquery"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel2(**kwargs)
+            self.rtree_index = get_cached_rtree(self.rtree_cache_dir, build_fn)
 
     def _read_bigquery(
         self,
@@ -833,9 +838,7 @@ class Sentinel2(DataSource):
             geometries: a list of geometries needed for each item
         """
         for item in items:
-            for suffix, band_names in self.BANDS:
-                if not is_raster_needed(band_names, self.config.band_sets):
-                    continue
+            for suffix, band_names in self.needed_bands:
                 if tile_store.is_raster_ready(item.name, band_names):
                     continue
 

rslearn/data_sources/google_earth_engine.py

@@ -20,7 +20,7 @@ from google.cloud import storage
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import DType, LayerConfig, RasterLayerConfig
+from rslearn.config import DType, LayerConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.dataset.materialize import RasterMaterializer
 from rslearn.dataset.window import Window
@@ -36,7 +36,7 @@ from rslearn.utils.raster_format import (
 )
 from rslearn.utils.rtree_index import RtreeIndex, get_cached_rtree
 
-from .data_source import DataSource, Item, QueryConfig
+from .data_source import DataSource, DataSourceContext, Item, QueryConfig
 
 logger = get_logger(__name__)
 
@@ -61,34 +61,55 @@ class GEE(DataSource, TileStore):
         self,
         collection_name: str,
         gcs_bucket_name: str,
-        bands: list[str],
-        index_cache_dir: UPath,
+        index_cache_dir: str,
         service_account_name: str,
         service_account_credentials: str,
+        bands: list[str] | None = None,
         filters: list[tuple[str, Any]] | None = None,
         dtype: DType | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new GEE instance.
 
         Args:
             collection_name: the Earth Engine ImageCollection to ingest images from
             gcs_bucket_name: the Cloud Storage bucket to export GEE images to
-            bands: the list of bands to ingest
             index_cache_dir: cache directory to store rtree index
             service_account_name: name of the service account to use for authentication
             service_account_credentials: service account credentials filename
+            bands: the list of bands to ingest, in case the layer config is not present
+                in the context.
             filters: optional list of tuples (property_name, property_value) to filter
                 images (using ee.Filter.eq)
             dtype: optional desired array data type. If the data obtained from GEE does
                 not match this type, then it is converted.
+            context: the data source context.
         """
         self.collection_name = collection_name
         self.gcs_bucket_name = gcs_bucket_name
-        self.bands = bands
-        self.index_cache_dir = index_cache_dir
         self.filters = filters
         self.dtype = dtype
 
+        # Get index cache dir depending on dataset path.
+        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)
+
+        # Get bands we need to export.
+        if context.layer_config is not None:
+            self.bands = [
+                band
+                for band_set in context.layer_config.band_sets
+                for band in band_set.bands
+            ]
+        elif bands is not None:
+            self.bands = bands
+        else:
+            raise ValueError(
+                "bands must be specified if layer_config is not present in the context"
+            )
+
         self.bucket = storage.Client().bucket(self.gcs_bucket_name)
 
         credentials = ee.ServiceAccountCredentials(
@@ -99,27 +120,6 @@ class GEE(DataSource, TileStore):
         self.index_cache_dir.mkdir(parents=True, exist_ok=True)
         self.rtree_index = get_cached_rtree(self.index_cache_dir, self._build_index)
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "GEE":
-        """Creates a new GEE instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("data_source is required in config")
-        d = config.data_source.config_dict
-        bands = [band for band_set in config.band_sets for band in band_set.bands]
-        kwargs = {
-            "collection_name": d["collection_name"],
-            "gcs_bucket_name": d["gcs_bucket_name"],
-            "bands": bands,
-            "service_account_name": d["service_account_name"],
-            "service_account_credentials": d["service_account_credentials"],
-            "filters": d.get("filters"),
-            "index_cache_dir": join_upath(ds_path, d["index_cache_dir"]),
-        }
-        if "dtype" in d:
-            kwargs["dtype"] = DType(d["dtype"])
-
-        return GEE(**kwargs)
-
     def get_collection(self) -> ee.ImageCollection:
         """Returns the Earth Engine image collection for this data source."""
         image_collection = ee.ImageCollection(self.collection_name)
@@ -578,7 +578,6 @@ class GEE(DataSource, TileStore):
             layer_name: the name of this layer
             layer_cfg: the config of this layer
         """
-        assert isinstance(layer_cfg, RasterLayerConfig)
         RasterMaterializer().materialize(
             TileStoreWithLayer(self, layer_name),
             window,
@@ -600,9 +599,10 @@ class GoogleSatelliteEmbeddings(GEE):
     def __init__(
         self,
         gcs_bucket_name: str,
-        index_cache_dir: UPath,
+        index_cache_dir: str,
         service_account_name: str,
         service_account_credentials: str,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Create a new GoogleSatelliteEmbeddings. See GEE for the arguments."""
         super().__init__(
@@ -612,22 +612,9 @@ class GoogleSatelliteEmbeddings(GEE):
             index_cache_dir=index_cache_dir,
             service_account_name=service_account_name,
             service_account_credentials=service_account_credentials,
+            context=context,
         )
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "GEE":
-        """Creates a new GEE instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("data_source is required in config")
-        d = config.data_source.config_dict
-        kwargs = {
-            "gcs_bucket_name": d["gcs_bucket_name"],
-            "index_cache_dir": join_upath(ds_path, d["index_cache_dir"]),
-            "service_account_name": d["service_account_name"],
-            "service_account_credentials": d["service_account_credentials"],
-        }
-        return GoogleSatelliteEmbeddings(**kwargs)
-
     # Override to add conversion to uint16.
     def item_to_image(self, item: Item) -> ee.image.Image:
         """Get the Image corresponding to the Item."""