rslearn.data_sources.hf_srtm

hf_srtm

Global SRTM void-filled elevation data from USGS, mirrored on Hugging Face by AI2.

This module provides: 1. A bulk download utility to fetch SRTM data from USGS EarthExplorer via the M2M API. This can be used to initialize a mirror of the data. 2. A data source that will pull from the AI2 Hugging Face mirror.

The SRTM dataset in USGS EarthExplorer is "srtm_v2" which contains void-filled elevation data from the Shuttle Radar Topography Mission. The bulk download fetches the highest resolution available: 1 arc-second (~30m) in the US and 3 arc-second (~90m) globally.

See https://www.usgs.gov/centers/eros/science/usgs-eros-archive-digital-elevation-shuttle-radar-topography-mission-srtm for details.

SRTM

Bases: DataSource

Data source for SRTM elevation data from the AI2 Hugging Face mirror.

The data is split into 1x1-degree tiles, with filenames like: N05/SRTM1N05W163V2.tif (1 arc-second, ~30m resolution) N05/SRTM3N05W163V2.tif (3 arc-second, ~90m resolution)

SRTM1 (1 arc-second) is available for some regions (primarily US territories), while SRTM3 (3 arc-second) is available globally. By default, SRTM1 is preferred when available for higher resolution. Set always_use_3arcsecond=True to always use the lower resolution SRTM3 data for consistency.

Items from this data source do not come with a time range. The band name will match that specified in the band set, which should have a single band (e.g. "dem").

Source code in rslearn/data_sources/hf_srtm.py

class SRTM(DataSource):
    """Data source for SRTM elevation data from the AI2 Hugging Face mirror.

    The data is split into 1x1-degree tiles, with filenames like:
    N05/SRTM1N05W163V2.tif (1 arc-second, ~30m resolution)
    N05/SRTM3N05W163V2.tif (3 arc-second, ~90m resolution)

    SRTM1 (1 arc-second) is available for some regions (primarily US territories),
    while SRTM3 (3 arc-second) is available globally. By default, SRTM1 is preferred
    when available for higher resolution. Set always_use_3arcsecond=True to always
    use the lower resolution SRTM3 data for consistency.

    Items from this data source do not come with a time range. The band name will match
    that specified in the band set, which should have a single band (e.g. "dem").
    """

    BASE_URL = (
        "https://huggingface.co/datasets/allenai/srtm-global-void-filled/resolve/main/"
    )
    FILE_LIST_FILENAME = "file_list.json"
    FILENAME_SUFFIX = "V2.tif"

    def __init__(
        self,
        timeout: timedelta = timedelta(seconds=10),
        cache_dir: str | None = None,
        always_use_3arcsecond: bool = False,
        context: DataSourceContext = DataSourceContext(),
    ):
        """Initialize a new SRTM instance.

        Args:
            timeout: timeout for requests.
            cache_dir: optional directory to cache the file list.
            always_use_3arcsecond: if True, always use 3 arc-second (SRTM3) data even
                when 1 arc-second (SRTM1) is available. Defaults to False, which
                prefers SRTM1 for higher resolution when available.
            context: the data source context.
        """
        # Get band name from context if possible, falling back to "dem".
        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 = "dem"

        self.timeout = timeout
        self.session = requests.session()
        self.always_use_3arcsecond = always_use_3arcsecond

        # Set the cache path if a cache_dir is provided.
        self.file_list_cache_path: UPath | None = None
        if cache_dir is not None:
            if context.ds_path is not None:
                cache_root = join_upath(context.ds_path, cache_dir)
            else:
                cache_root = UPath(cache_dir)
            cache_root.mkdir(parents=True, exist_ok=True)
            self.file_list_cache_path = join_upath(cache_root, self.FILE_LIST_FILENAME)

        self._basename_to_item, self._tile_to_item = self._load_file_index()

    def _load_file_index(
        self,
    ) -> tuple[dict[str, Item], dict[tuple[int, int], Item]]:
        """Load the file list and build indices for lookups."""
        file_list = self._load_file_list_json()
        if not isinstance(file_list, list):
            raise ValueError("expected file_list.json to be a list of filenames")

        basename_to_item: dict[str, Item] = {}
        tile_to_item: dict[tuple[int, int], Item] = {}

        for entry in file_list:
            if not isinstance(entry, str):
                raise ValueError(
                    "expected file_list.json to contain only string filenames"
                )
            basename = os.path.basename(entry)

            # Check if this is SRTM1 or SRTM3 based on filename prefix
            is_srtm1 = basename.startswith("SRTM1")

            # Skip SRTM1 files if always_use_3arcsecond is enabled
            if self.always_use_3arcsecond and is_srtm1:
                continue

            lat_min, lon_min = self._parse_tile_basename(basename)
            geometry = STGeometry(
                WGS84_PROJECTION,
                shapely.box(lon_min, lat_min, lon_min + 1, lat_min + 1),
                None,
            )
            item = Item(entry, geometry)

            key = (lon_min, lat_min)

            # For tile_to_item, prefer SRTM1 over SRTM3 when not using always_use_3arcsecond
            if key in tile_to_item:
                existing_is_srtm1 = os.path.basename(tile_to_item[key].name).startswith(
                    "SRTM1"
                )
                # Only replace if current is SRTM1 and existing is not
                if is_srtm1 and not existing_is_srtm1:
                    tile_to_item[key] = item
                # Keep existing if it's SRTM1 and current is not
                elif existing_is_srtm1 and not is_srtm1:
                    pass
                else:
                    # Same type, keep the existing one
                    pass
            else:
                tile_to_item[key] = item

            if basename not in basename_to_item:
                basename_to_item[basename] = item

        logger.info(
            f"Loaded {len(tile_to_item)} SRTM tiles from Hugging Face file list"
        )
        return basename_to_item, tile_to_item

    def _load_file_list_json(self) -> list[Any]:
        """Load file list JSON, optionally from cache."""
        if self.file_list_cache_path is not None and self.file_list_cache_path.exists():
            with self.file_list_cache_path.open() as f:
                file_list = json.load(f)
            logger.info(f"Loaded SRTM file list cache from {self.file_list_cache_path}")
            return file_list

        response = self.session.get(
            self.BASE_URL + self.FILE_LIST_FILENAME,
            timeout=self.timeout.total_seconds(),
        )
        response.raise_for_status()
        file_list = response.json()
        if self.file_list_cache_path is not None:
            with open_atomic(self.file_list_cache_path, "w") as f:
                json.dump(file_list, f)
        return file_list

    def _parse_tile_basename(self, basename: str) -> tuple[int, int]:
        """Parse a tile basename into (lat_min, lon_min)."""
        match = re.match(
            r"^SRTM\d([NS])(\d{2})([EW])(\d{3})V2\.tif$",
            basename,
        )
        if match is None:
            raise ValueError(f"invalid SRTM tile filename: {basename}")

        lat_sign, lat_str, lon_sign, lon_str = match.groups()
        lat_degrees = int(lat_str)
        lon_degrees = int(lon_str)

        if lat_sign == "N":
            lat_min = lat_degrees
        elif lat_sign == "S":
            lat_min = -lat_degrees
        else:
            raise ValueError(f"invalid SRTM tile latitude for filename: {basename}")

        if lon_sign == "E":
            lon_min = lon_degrees
        elif lon_sign == "W":
            lon_min = -lon_degrees
        else:
            raise ValueError(f"invalid SRTM tile longitude for filename: {basename}")

        return lat_min, lon_min

    def get_item_by_name(self, name: str) -> Item:
        """Gets an item by name.

        Args:
            name: the name of the item to get. For SRTM, the item name is the filename
                of the GeoTIFF tile.

        Returns:
            the Item object
        """
        basename = os.path.basename(name)
        item = self._basename_to_item.get(basename)
        if item is not None:
            return item
        raise ValueError(f"unknown SRTM tile name: {name}")

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[Item]]]:
        """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.
        """
        # It only makes sense to create mosaic from the SRTM data since it is one
        # global layer (but spatially split up into items).
        if query_config.space_mode != SpaceMode.MOSAIC or query_config.max_matches != 1:
            raise ValueError(
                "expected mosaic with max_matches=1 for the query configuration"
            )
        if query_config.min_matches != 0:
            raise ValueError("min_matches is not supported for SRTM; set min_matches=0")

        groups = []
        for geometry in geometries:
            wgs84_geometry = geometry.to_wgs84()
            items = []
            # Use a set to avoid adding the same item multiple times for different
            # components of a MultiPolygon or geometry collection.
            seen_tiles = set()
            for shp in flatten_shape(wgs84_geometry.shp):
                shp_bounds = shp.bounds
                cell_bounds = (
                    math.floor(shp_bounds[0]),
                    math.floor(shp_bounds[1]),
                    math.ceil(shp_bounds[2]),
                    math.ceil(shp_bounds[3]),
                )
                for lon_min in range(cell_bounds[0], cell_bounds[2]):
                    for lat_min in range(cell_bounds[1], cell_bounds[3]):
                        if (lon_min, lat_min) not in self._tile_to_item:
                            continue
                        if (lon_min, lat_min) in seen_tiles:
                            continue
                        items.append(self._tile_to_item[(lon_min, lat_min)])
                        seen_tiles.add((lon_min, lat_min))

            logger.debug(f"Got {len(items)} items (grid cells) for geometry")
            groups.append([MatchedItemGroup(items, geometry.time_range)])

        return groups

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

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[Item],
        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:
            if tile_store.is_raster_ready(item, [self.band_name]):
                continue

            url = self.BASE_URL + item.name
            logger.debug(f"Downloading SRTM data for {item.name} from {url}")
            response = self.session.get(
                url, stream=True, timeout=self.timeout.total_seconds()
            )

            if response.status_code == 404:
                logger.warning(
                    f"Skipping item {item.name} because there is no data at that cell"
                )
                continue
            response.raise_for_status()

            with tempfile.TemporaryDirectory() as tmp_dir:
                local_fname = os.path.join(tmp_dir, "data.tif")
                with open(local_fname, "wb") as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)

                logger.debug(f"Ingesting data for {item.name}")
                tile_store.write_raster_file(
                    item,
                    [self.band_name],
                    UPath(local_fname),
                    time_range=item.geometry.time_range,
                )

get_item_by_name

get_item_by_name(name: str) -> Item

Gets an item by name.

Parameters:

Name	Type	Description	Default
name	str	the name of the item to get. For SRTM, the item name is the filename of the GeoTIFF tile.	required

Returns:

Type	Description
Item	the Item object

Source code in rslearn/data_sources/hf_srtm.py

def get_item_by_name(self, name: str) -> Item:
    """Gets an item by name.

    Args:
        name: the name of the item to get. For SRTM, the item name is the filename
            of the GeoTIFF tile.

    Returns:
        the Item object
    """
    basename = os.path.basename(name)
    item = self._basename_to_item.get(basename)
    if item is not None:
        return item
    raise ValueError(f"unknown SRTM tile name: {name}")

get_items

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

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[Item]]]	List of groups of items that should be retrieved for each geometry.

Source code in rslearn/data_sources/hf_srtm.py

def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[Item]]]:
    """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.
    """
    # It only makes sense to create mosaic from the SRTM data since it is one
    # global layer (but spatially split up into items).
    if query_config.space_mode != SpaceMode.MOSAIC or query_config.max_matches != 1:
        raise ValueError(
            "expected mosaic with max_matches=1 for the query configuration"
        )
    if query_config.min_matches != 0:
        raise ValueError("min_matches is not supported for SRTM; set min_matches=0")

    groups = []
    for geometry in geometries:
        wgs84_geometry = geometry.to_wgs84()
        items = []
        # Use a set to avoid adding the same item multiple times for different
        # components of a MultiPolygon or geometry collection.
        seen_tiles = set()
        for shp in flatten_shape(wgs84_geometry.shp):
            shp_bounds = shp.bounds
            cell_bounds = (
                math.floor(shp_bounds[0]),
                math.floor(shp_bounds[1]),
                math.ceil(shp_bounds[2]),
                math.ceil(shp_bounds[3]),
            )
            for lon_min in range(cell_bounds[0], cell_bounds[2]):
                for lat_min in range(cell_bounds[1], cell_bounds[3]):
                    if (lon_min, lat_min) not in self._tile_to_item:
                        continue
                    if (lon_min, lat_min) in seen_tiles:
                        continue
                    items.append(self._tile_to_item[(lon_min, lat_min)])
                    seen_tiles.add((lon_min, lat_min))

        logger.debug(f"Got {len(items)} items (grid cells) for geometry")
        groups.append([MatchedItemGroup(items, geometry.time_range)])

    return groups

deserialize_item

deserialize_item(serialized_item: dict) -> Item

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/hf_srtm.py

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

ingest

ingest(tile_store: TileStoreWithLayer, items: list[Item], 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[Item]	the items to ingest	required
geometries	list[list[STGeometry]]	a list of geometries needed for each item	required

Source code in rslearn/data_sources/hf_srtm.py

def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[Item],
    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:
        if tile_store.is_raster_ready(item, [self.band_name]):
            continue

        url = self.BASE_URL + item.name
        logger.debug(f"Downloading SRTM data for {item.name} from {url}")
        response = self.session.get(
            url, stream=True, timeout=self.timeout.total_seconds()
        )

        if response.status_code == 404:
            logger.warning(
                f"Skipping item {item.name} because there is no data at that cell"
            )
            continue
        response.raise_for_status()

        with tempfile.TemporaryDirectory() as tmp_dir:
            local_fname = os.path.join(tmp_dir, "data.tif")
            with open(local_fname, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)

            logger.debug(f"Ingesting data for {item.name}")
            tile_store.write_raster_file(
                item,
                [self.band_name],
                UPath(local_fname),
                time_range=item.geometry.time_range,
            )

bulk_download_srtm

bulk_download_srtm(output_dir: str, num_workers: int = 4, timeout: timedelta = timedelta(minutes=5)) -> None

Bulk download SRTM data from USGS EarthExplorer.

Downloads all SRTM tiles to the specified output directory. Uses atomic rename to ensure partially downloaded files are not included. Files that already exist in the output directory are skipped.

The scene list is cached in scenes.json in the output directory to avoid re-querying on subsequent runs.

Requires M2M_USERNAME and M2M_TOKEN environment variables to be set.

Parameters:

Name	Type	Description	Default
output_dir	str	directory to save downloaded files	required
num_workers	int	number of parallel download workers	4
timeout	timedelta	timeout for API requests and downloads	timedelta(minutes=5)

Source code in rslearn/data_sources/hf_srtm.py

def bulk_download_srtm(
    output_dir: str,
    num_workers: int = 4,
    timeout: timedelta = timedelta(minutes=5),
) -> None:
    """Bulk download SRTM data from USGS EarthExplorer.

    Downloads all SRTM tiles to the specified output directory. Uses atomic
    rename to ensure partially downloaded files are not included. Files that
    already exist in the output directory are skipped.

    The scene list is cached in scenes.json in the output directory to avoid
    re-querying on subsequent runs.

    Requires M2M_USERNAME and M2M_TOKEN environment variables to be set.

    Args:
        output_dir: directory to save downloaded files
        num_workers: number of parallel download workers
        timeout: timeout for API requests and downloads
    """
    os.makedirs(output_dir, exist_ok=True)

    client = _get_cached_m2m_client(timeout)
    scenes = _fetch_all_scenes(client, output_dir)

    # Filter out scenes that are already downloaded
    download_tasks = []
    skipped = 0
    for display_id, scene in scenes.items():
        # Use display_id as the filename with .tif extension for GeoTIFF
        output_path = os.path.join(output_dir, f"{display_id}.tif")
        if os.path.exists(output_path):
            logger.debug(f"Skipping {display_id} - already downloaded")
            skipped += 1
        else:
            download_tasks.append(
                {
                    "scene": scene,
                    "output_path": output_path,
                    "timeout": timeout,
                }
            )

    logger.info(
        f"Need to download {len(download_tasks)} scenes ({skipped} already downloaded)"
    )

    if not download_tasks:
        logger.info("All scenes already downloaded!")
        return

    # Download in parallel using multiprocessing.Pool
    # Each worker creates its own M2M API client to get download URLs
    logger.info(f"Starting downloads with {num_workers} workers...")
    with make_pool_and_star_imap_unordered(
        num_workers, _worker_download_scene, download_tasks
    ) as outputs:
        for output_path in outputs:
            logger.info(f"Downloaded {output_path}")

main

main() -> None

Command-line entry point for bulk SRTM download.

Requires M2M_USERNAME and M2M_TOKEN environment variables to be set.

Source code in rslearn/data_sources/hf_srtm.py

def main() -> None:
    """Command-line entry point for bulk SRTM download.

    Requires M2M_USERNAME and M2M_TOKEN environment variables to be set.
    """
    import argparse
    import logging

    parser = argparse.ArgumentParser(
        description="Bulk download SRTM data from USGS EarthExplorer. "
        "Requires M2M_USERNAME and M2M_TOKEN environment variables."
    )
    parser.add_argument(
        "output_dir",
        help="Directory to save downloaded SRTM files",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Number of parallel download workers (default: 4)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=30,
        help="Timeout in seconds for API requests and downloads (default: 30)",
    )

    args = parser.parse_args()

    # Configure logging based on RSLEARN_LOGLEVEL
    log_level = os.environ.get("RSLEARN_LOGLEVEL", "INFO")
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s %(levelname)-6s %(name)s %(message)s",
    )
    # Enable urllib3 logging to see retry information at DEBUG level
    logging.getLogger("urllib3").setLevel(log_level)

    bulk_download_srtm(
        output_dir=args.output_dir,
        num_workers=args.workers,
        timeout=timedelta(seconds=args.timeout),
    )