rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/data_sources/worldpop.py

@@ -0,0 +1,153 @@
+"""Data from worldpop.org."""
+
+import random
+from datetime import timedelta
+from html.parser import HTMLParser
+from urllib.parse import urljoin
+
+import requests
+from upath import UPath
+
+from rslearn.config import LayerType
+from rslearn.data_sources import DataSourceContext
+from rslearn.data_sources.local_files import LocalFiles
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import join_upath, open_atomic
+
+logger = get_logger(__name__)
+
+
+class LinkExtractor(HTMLParser):
+    """Extract links from HTML.
+
+    The links attribute will be filled with the href attribute of all links that appear
+    on the HTML page.
+    """
+
+    def __init__(self) -> None:
+        """Create a new LinkExtractor."""
+        super().__init__()
+        self.links: list[str] = []
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        """Handle start of tag from the HTML parsing."""
+        if tag.lower() != "a":
+            return
+        for name, value in attrs:
+            if name.lower() != "href":
+                continue
+            if value is None:
+                continue
+            self.links.append(value)
+
+
+class WorldPop(LocalFiles):
+    """World population data from worldpop.org.
+
+    Currently, this only supports the WorldPop Constrained 2020 100 m Resolution
+    dataset. See https://hub.worldpop.org/project/categories?id=3 for details.
+
+    The data is split by country. We implement with LocalFiles data source for
+    simplicity, but it means that all of the data must be downloaded first.
+    """
+
+    INDEX_URLS = [
+        "https://data.worldpop.org/GIS/Population/Global_2000_2020_Constrained/2020/BSGM/",
+        "https://data.worldpop.org/GIS/Population/Global_2000_2020_Constrained/2020/maxar_v1/",
+    ]
+    FILENAME_SUFFIX = "_ppp_2020_constrained.tif"
+
+    def __init__(
+        self,
+        worldpop_dir: str,
+        timeout: timedelta = timedelta(seconds=30),
+        context: DataSourceContext = DataSourceContext(),
+    ):
+        """Create a new WorldPop.
+
+        Args:
+            worldpop_dir: the directory to extract the WorldPop GeoTIFF files. For
+                high performance, this should be a local directory; if the dataset is
+                remote, prefix with a protocol ("file://") to use a local directory
+                instead of a path relative to the dataset path.
+            timeout: timeout for HTTP requests.
+            context: the data source context.
+        """
+        if context.ds_path is not None:
+            worldpop_upath = join_upath(context.ds_path, worldpop_dir)
+        else:
+            worldpop_upath = UPath(worldpop_dir)
+        worldpop_upath.mkdir(parents=True, exist_ok=True)
+        self.download_worldpop_data(worldpop_upath, timeout)
+        super().__init__(
+            src_dir=worldpop_upath,
+            layer_type=LayerType.RASTER,
+            context=context,
+        )
+
+    def download_worldpop_data(self, worldpop_dir: UPath, timeout: timedelta) -> None:
+        """Download and extract the WorldPop data.
+
+        If the data was previously downloaded, this function returns quickly.
+
+        Args:
+            worldpop_dir: the directory to download to.
+            timeout: timeout for HTTP requests.
+        """
+        completed_fname = worldpop_dir / "completed"
+        if completed_fname.exists():
+            return
+
+        # Scan the index URLs to get all the per-country subfolders.
+        # These should be four characters with slash at the end, like "USA/".
+        country_urls = []
+        for index_url in self.INDEX_URLS:
+            logger.info(f"Getting per-country subfolders from {index_url}")
+            response = requests.get(index_url, timeout=timeout.total_seconds())
+            response.raise_for_status()
+            parser = LinkExtractor()
+            parser.feed(response.text)
+            country_urls.extend(
+                [
+                    urljoin(index_url, href)
+                    for href in parser.links
+                    if len(href) == 4 and href[3] == "/"
+                ]
+            )
+
+        logger.info(f"Got {len(country_urls)} country subfolders to download")
+        # Shuffling here enables the user to run multiple processes to speed up the
+        # download.
+        random.shuffle(country_urls)
+
+        # Now iterate over the country-level URLs and download the GeoTIFF.
+        for country_url in country_urls:
+            response = requests.get(country_url, timeout=timeout.total_seconds())
+            response.raise_for_status()
+            parser = LinkExtractor()
+            parser.feed(response.text)
+            tif_links = [
+                urljoin(country_url, href)
+                for href in parser.links
+                if href.endswith(self.FILENAME_SUFFIX)
+            ]
+            if len(tif_links) != 1:
+                raise ValueError(
+                    f"expected {country_url} to contain one GeoTIFF ending in {self.FILENAME_SUFFIX} but got {parser.links}"
+                )
+
+            country_fname = tif_links[0].split("/")[-1]
+            dst_fname = worldpop_dir / country_fname
+            if dst_fname.exists():
+                continue
+
+            logger.info(f"Downloading from {tif_links[0]} to {dst_fname}")
+            with requests.get(
+                tif_links[0], stream=True, timeout=timeout.total_seconds()
+            ) as r:
+                r.raise_for_status()
+                with open_atomic(dst_fname, "wb") as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+
+        completed_fname.touch()

rslearn/data_sources/xyz_tiles.py

@@ -13,15 +13,17 @@ import rasterio.warp
 import shapely
 from PIL import Image
 from rasterio.crs import CRS
-from upath import UPath
+from rasterio.enums import Resampling
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import LayerConfig, QueryConfig
 from rslearn.dataset import Window
+from rslearn.dataset.materialize import RasterMaterializer
+from rslearn.tile_stores import TileStore, TileStoreWithLayer
 from rslearn.utils import PixelBounds, Projection, STGeometry
 from rslearn.utils.array import copy_spatial_array
+from rslearn.utils.raster_format import get_transform_from_projection_and_bounds
 
-from .data_source import DataSource, Item
-from .raster_source import ArrayWithTransform, materialize_raster
+from .data_source import DataSource, DataSourceContext, Item
 from .utils import match_candidate_items_to_window
 
 WEB_MERCATOR_EPSG = 3857
@@ -81,58 +83,24 @@ def read_from_tile_callback(
     return data
 
 
-class XyzItem(Item):
-    """An item in the XyzTiles data source.
-
-    Each item represents one layer of tiles. Often there is only one itm in the data
-    source, but if there are multiple then they should correspond to different time
-    ranges.
-    """
-
-    def __init__(self, name: str, geometry: STGeometry, url_template: str):
-        """Creates a new XyzItem.
-
-        Args:
-            name: unique name of the item
-            geometry: the spatial and temporal extent of the item
-            url_template: the URL template for an xyz tile.
-        """
-        super().__init__(name, geometry)
-        self.url_template = url_template
-
-    def serialize(self) -> dict:
-        """Serializes the item to a JSON-encodable dictionary."""
-        d = super().serialize()
-        d["url_template"] = self.url_template
-        return d
-
-    @staticmethod
-    def deserialize(d: dict) -> Item:
-        """Deserializes an item from a JSON-decoded dictionary."""
-        item = super(XyzItem, XyzItem).deserialize(d)
-        return XyzItem(
-            name=item.name, geometry=item.geometry, url_template=d["url_template"]
-        )
-
-
-class XyzTiles(DataSource):
+class XyzTiles(DataSource, TileStore):
     """A data source for web xyz image tiles.
 
     These tiles are usually in WebMercator projection, but different CRS can be
     configured here.
     """
 
-    item_name = "xyz_tiles"
-
     def __init__(
         self,
         url_templates: list[str],
         time_ranges: list[tuple[datetime, datetime]],
         zoom: int,
-        crs: CRS = CRS.from_epsg(WEB_MERCATOR_EPSG),
+        crs: str | CRS = CRS.from_epsg(WEB_MERCATOR_EPSG),
         total_units: float = WEB_MERCATOR_UNITS,
         offset: float = WEB_MERCATOR_UNITS / 2,
         tile_size: int = 256,
+        band_names: list[str] = ["R", "G", "B"],
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize an XyzTiles instance.
 
@@ -152,14 +120,22 @@ class XyzTiles(DataSource):
                 the pixel size to map from projection coordinates to pixel coordinates.
             offset: offset added to projection units when converting to tile positions.
             tile_size: size in pixels of each tile. Tiles must be square.
+            band_names: what to name the bands that we read.
+            context: the data source context.
         """
         self.url_templates = url_templates
         self.time_ranges = time_ranges
         self.zoom = zoom
-        self.crs = crs
         self.total_units = total_units
         self.offset = offset
         self.tile_size = tile_size
+        self.band_names = band_names
+
+        # Convert to CRS if needed.
+        if isinstance(crs, str):
+            self.crs = CRS.from_string(crs)
+        else:
+            self.crs = crs
 
         # Compute total number of pixels (a function of the zoom level and tile size).
         self.total_pixels = tile_size * (2**zoom)
@@ -169,7 +145,7 @@ class XyzTiles(DataSource):
         self.pixel_offset = int(self.offset / self.pixel_size)
         # Compute the extent in pixel coordinates as an STGeometry.
         # Note that pixel coordinates are prior to applying the offset.
-        shp = shapely.box(
+        self.shp = shapely.box(
             -self.total_pixels // 2,
             -self.total_pixels // 2,
             self.total_pixels // 2,
@@ -179,32 +155,10 @@ class XyzTiles(DataSource):
 
         self.items = []
         for url_template, time_range in zip(self.url_templates, self.time_ranges):
-            geometry = STGeometry(self.projection, shp, time_range)
-            item = XyzItem(self.item_name, geometry, url_template)
+            geometry = STGeometry(self.projection, self.shp, time_range)
+            item = Item(url_template, geometry)
             self.items.append(item)
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "XyzTiles":
-        """Creates a new XyzTiles instance from a configuration dictionary."""
-        d = config.data_source.config_dict
-        time_ranges = []
-        for str1, str2 in d["time_ranges"]:
-            time1 = datetime.fromisoformat(str1)
-            time2 = datetime.fromisoformat(str2)
-            time_ranges.append((time1, time2))
-        kwargs = dict(
-            url_templates=d["url_templates"], zoom=d["zoom"], time_ranges=time_ranges
-        )
-        if "crs" in d:
-            kwargs["crs"] = CRS.from_string(d["crs"])
-        if "total_units" in d:
-            kwargs["total_units"] = d["total_units"]
-        if "offset" in d:
-            kwargs["offset"] = d["offset"]
-        if "tile_size" in d:
-            kwargs["tile_size"] = d["tile_size"]
-        return XyzTiles(**kwargs)
-
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
     ) -> list[list[list[Item]]]:
@@ -232,7 +186,7 @@ class XyzTiles(DataSource):
 
     def deserialize_item(self, serialized_item: Any) -> Item:
         """Deserializes an item from JSON-decoded data."""
-        return XyzItem.deserialize(serialized_item)
+        return Item.deserialize(serialized_item)
 
     def read_tile(self, url_template: str, col: int, row: int) -> npt.NDArray[Any]:
         """Read the tile at specified column and row.
@@ -249,8 +203,11 @@ class XyzTiles(DataSource):
         url = url.replace("{x}", str(col))
         url = url.replace("{y}", str(row))
         url = url.replace("{z}", str(self.zoom))
-        image = Image.open(urllib.request.urlopen(url))
-        return np.array(image).transpose(2, 0, 1)
+        image = np.array(Image.open(urllib.request.urlopen(url)))
+        # Handle grayscale images (add single-band channel dimension).
+        if len(image.shape) == 2:
+            image = image[:, :, None]
+        return image.transpose(2, 0, 1)
 
     def read_bounds(self, url_template: str, bounds: PixelBounds) -> npt.NDArray[Any]:
         """Reads the portion of the raster in the specified bounds.
@@ -275,6 +232,122 @@ class XyzTiles(DataSource):
             self.tile_size,
         )
 
+    def is_raster_ready(
+        self, layer_name: str, item_name: str, bands: list[str]
+    ) -> bool:
+        """Checks if this raster has been written to the store.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+            bands: the list of bands identifying which specific raster to read.
+
+        Returns:
+            whether there is a raster in the store matching the source, item, and
+                bands.
+        """
+        # Always ready since we wrap accesses to the XYZ tile URL.
+        return True
+
+    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
+        """Get the sets of bands that have been stored for the specified item.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+
+        Returns:
+            a list of lists of bands that are in the tile store (with one raster
+                stored corresponding to each inner list). If no rasters are ready for
+                this item, returns empty list.
+        """
+        return [self.band_names]
+
+    def get_raster_bounds(
+        self, layer_name: str, item_name: str, bands: list[str], projection: Projection
+    ) -> PixelBounds:
+        """Get the bounds of the raster in the specified projection.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item to check.
+            bands: the list of bands identifying which specific raster to read. These
+                bands must match the bands of a stored raster.
+            projection: the projection to get the raster's bounds in.
+
+        Returns:
+            the bounds of the raster in the projection.
+        """
+        geom = STGeometry(self.projection, self.shp, None).to_projection(projection)
+        return (
+            int(geom.shp.bounds[0]),
+            int(geom.shp.bounds[1]),
+            int(geom.shp.bounds[2]),
+            int(geom.shp.bounds[3]),
+        )
+
+    def read_raster(
+        self,
+        layer_name: str,
+        item_name: str,
+        bands: list[str],
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling = Resampling.bilinear,
+    ) -> npt.NDArray[Any]:
+        """Read raster data from the store.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item to read.
+            bands: the list of bands identifying which specific raster to read. These
+                bands must match the bands of a stored raster.
+            projection: the projection to read in.
+            bounds: the bounds to read.
+            resampling: the resampling method to use in case reprojection is needed.
+
+        Returns:
+            the raster data
+        """
+        # Validate bands.
+        if bands != self.band_names:
+            raise ValueError(
+                f"expected request for bands {self.band_names} but requested {bands}"
+            )
+
+        # Read a raster matching the given bounds but projected onto the projection of
+        # the xyz tiles.
+        request_geometry = STGeometry(projection, shapely.box(*bounds), None)
+        projected_geometry = request_geometry.to_projection(self.projection)
+        projected_bounds = (
+            math.floor(projected_geometry.shp.bounds[0]),
+            math.floor(projected_geometry.shp.bounds[1]),
+            math.ceil(projected_geometry.shp.bounds[2]),
+            math.ceil(projected_geometry.shp.bounds[3]),
+        )
+        # The item name is the URL template.
+        url_template = item_name
+        array = self.read_bounds(url_template, projected_bounds)
+        # Now project it back to the requested geometry.
+        src_transform = get_transform_from_projection_and_bounds(
+            self.projection, projected_bounds
+        )
+        dst_transform = get_transform_from_projection_and_bounds(projection, bounds)
+        dst_array = np.zeros(
+            (array.shape[0], bounds[3] - bounds[1], bounds[2] - bounds[0]),
+            dtype=array.dtype,
+        )
+        rasterio.warp.reproject(
+            source=array,
+            src_crs=self.projection.crs,
+            src_transform=src_transform,
+            destination=dst_array,
+            dst_crs=projection.crs,
+            dst_transform=dst_transform,
+            resampling=resampling,
+        )
+        return dst_array
+
     def materialize(
         self,
         window: Window,
@@ -290,40 +363,10 @@ class XyzTiles(DataSource):
             layer_name: the name of this layer
             layer_cfg: the config of this layer
         """
-        assert len(item_groups) == 1 and len(item_groups[0]) == 1
-        item = item_groups[0][0]
-        assert isinstance(item, XyzItem)
-
-        # Read a raster matching the bounds of the window's bounds projected onto the
-        # projection of the xyz tiles.
-        assert isinstance(layer_cfg, RasterLayerConfig)
-        band_cfg = layer_cfg.band_sets[0]
-        window_projection, window_bounds = band_cfg.get_final_projection_and_bounds(
-            window.projection, window.bounds
-        )
-        window_geometry = STGeometry(
-            window_projection, shapely.box(*window_bounds), None
+        RasterMaterializer().materialize(
+            TileStoreWithLayer(self, layer_name),
+            window,
+            layer_name,
+            layer_cfg,
+            item_groups,
         )
-        projected_geometry = window_geometry.to_projection(self.projection)
-        projected_bounds = [
-            math.floor(projected_geometry.shp.bounds[0]),
-            math.floor(projected_geometry.shp.bounds[1]),
-            math.ceil(projected_geometry.shp.bounds[2]),
-            math.ceil(projected_geometry.shp.bounds[3]),
-        ]
-        projected_raster = self.read_bounds(item.url_template, projected_bounds)
-
-        # Attach the transform to the raster.
-        src_transform = rasterio.transform.Affine(
-            self.projection.x_resolution,
-            0,
-            projected_bounds[0] * self.projection.x_resolution,
-            0,
-            self.projection.y_resolution,
-            projected_bounds[1] * self.projection.y_resolution,
-        )
-        array_with_transform = ArrayWithTransform(
-            projected_raster, self.projection.crs, src_transform
-        )
-
-        materialize_raster(array_with_transform, window, layer_name, band_cfg)

rslearn/dataset/__init__.py

@@ -1,6 +1,12 @@
 """rslearn dataset storage and operations."""
 
 from .dataset import Dataset
-from .window import Window, WindowLayerData
+from .window import Window, WindowLayerData, get_window_layer_dir, get_window_raster_dir
 
-__all__ = ("Dataset", "Window", "WindowLayerData")
+__all__ = (
+    "Dataset",
+    "Window",
+    "WindowLayerData",
+    "get_window_layer_dir",
+    "get_window_raster_dir",
+)

rslearn/dataset/add_windows.py

@@ -25,7 +25,7 @@ def add_windows_from_geometries(
     window_size: int | None = None,
     time_range: tuple[datetime, datetime] | None = None,
     use_utm: bool = False,
-):
+) -> list[Window]:
     """Create windows based on a list of STGeometry.
 
     Args:
@@ -131,7 +131,7 @@ def add_windows_from_geometries(
                     f"_{time_range[0].isoformat()}_{time_range[1].isoformat()}"
                 )
         window = Window(
-            path=dataset.path / "windows" / group / cur_window_name,
+            storage=dataset.storage,
             group=group,
             name=cur_window_name,
             projection=cur_projection,

rslearn/dataset/dataset.py

@@ -1,16 +1,19 @@
 """rslearn dataset class."""
 
 import json
-import multiprocessing
+from typing import Any
 
-import tqdm
 from upath import UPath
 
-from rslearn.config import TileStoreConfig, load_layer_config
+from rslearn.config import DatasetConfig
+from rslearn.log_utils import get_logger
+from rslearn.template_params import substitute_env_vars_in_string
 from rslearn.tile_stores import TileStore, load_tile_store
 
 from .window import Window
 
+logger = get_logger(__name__)
+
 
 class Dataset:
     """A rslearn dataset.
@@ -20,7 +23,7 @@ class Dataset:
     .. code-block:: none
 
         dataset/
-            config.json
+            config.json  # optional, if config provided as runtime object
             windows/
                 group1/
                     epsg:3857_10_623565_1528020/
@@ -37,72 +40,58 @@ class Dataset:
     materialize.
     """
 
-    def __init__(self, path: UPath) -> None:
+    def __init__(
+        self,
+        path: UPath,
+        disabled_layers: list[str] = [],
+        dataset_config: DatasetConfig | None = None,
+    ) -> None:
         """Initializes a new Dataset.
 
         Args:
             path: the root directory of the dataset
+            disabled_layers: list of layers to disable
+            dataset_config: optional dataset configuration to use instead of loading from the dataset directory
         """
         self.path = path
 
-        # Load dataset configuration.
-        with (self.path / "config.json").open("r") as f:
-            config = json.load(f)
-            self.layers = {
-                layer_name: load_layer_config(d)
-                for layer_name, d in config["layers"].items()
-            }
-            self.tile_store_config = TileStoreConfig.from_config(config["tile_store"])
-            self.materializer_name = config.get("materialize")
+        if dataset_config is None:
+            # Load dataset configuration from the dataset directory.
+            with (self.path / "config.json").open("r") as f:
+                config_content = f.read()
+                config_content = substitute_env_vars_in_string(config_content)
+                dataset_config = DatasetConfig.model_validate(
+                    json.loads(config_content)
+                )
+
+        self.layers = {}
+        for layer_name, layer_config in dataset_config.layers.items():
+            if layer_name in disabled_layers:
+                logger.warning(f"Layer {layer_name} is disabled")
+                continue
+            self.layers[layer_name] = layer_config
+
+        self.tile_store_config = dataset_config.tile_store
+        self.storage = (
+            dataset_config.storage.instantiate_window_storage_factory().get_storage(
+                self.path
+            )
+        )
 
     def load_windows(
         self,
         groups: list[str] | None = None,
         names: list[str] | None = None,
-        show_progress: bool = False,
-        workers: int = 0,
+        **kwargs: Any,
     ) -> list[Window]:
         """Load the windows in the dataset.
 
         Args:
             groups: an optional list of groups to filter loading
             names: an optional list of window names to filter loading
-            show_progress: whether to show tqdm progress bar
-            workers: number of parallel workers, default 0 (use main thread only to load windows)
+            kwargs: optional keyword arguments to pass to WindowStorage.get_windows.
         """
-        window_dirs = []
-        if not groups:
-            groups = []
-            for p in (self.path / "windows").iterdir():
-                groups.append(p.name)
-        for group in groups:
-            group_dir = self.path / "windows" / group
-            if names:
-                cur_names = names
-            else:
-                cur_names = []
-                for p in group_dir.iterdir():
-                    cur_names.append(p.name)
-
-            for window_name in cur_names:
-                window_dir = group_dir / window_name
-                window_dirs.append(window_dir)
-
-        if workers == 0:
-            windows = [Window.load(window_dir) for window_dir in window_dirs]
-        else:
-            p = multiprocessing.Pool(workers)
-            outputs = p.imap_unordered(Window.load, window_dirs)
-            if show_progress:
-                outputs = tqdm.tqdm(
-                    outputs, total=len(window_dirs), desc="Loading windows"
-                )
-            windows = []
-            for window in outputs:
-                windows.append(window)
-            p.close()
-
-        return windows
+        return self.storage.get_windows(groups=groups, names=names, **kwargs)
 
     def get_tile_store(self) -> TileStore:
         """Get the tile store associated with this dataset.