diff --git a/iterable_example.ipynb b/iterable_example.ipynb new file mode 100644 index 00000000..8dc7e64d --- /dev/null +++ b/iterable_example.ipynb @@ -0,0 +1,86 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "760550cf-4feb-44b4-9f7f-ce065fde279b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[-93.01025390625037, 44.99588261816585, -92.98828125000037, 45.01141864227766]\n", + "https://planetarycomputer.microsoft.com/api/stac/v1/search?limit=10&bbox=-93.01025390625037%2C44.99588261816585%2C-92.98828125000037%2C45.01141864227766&collections=naip&fields=\n", + "4 4 4\n", + "4\n", + "CPU times: user 383 ms, sys: 7.44 ms, total: 391 ms\n", + "Wall time: 1.08 s\n" + ] + } + ], + "source": [ + "%%time\n", + "from pystac_client import Client\n", + "from rio_tiler.mosaic import mosaic_reader\n", + "from rio_tiler.io import Reader\n", + "from rio_tiler.models import ImageData\n", + "import morecantile\n", + "\n", + "tms = morecantile.tms.get(\"WebMercatorQuad\")\n", + "x, y, z = tms.tile(-93,45,14)\n", + "bbox = list(tms.bounds(morecantile.Tile(x, y, z)))\n", + "print(bbox)\n", + "\n", + "def reader(asset: str, *args, **kwargs) -> ImageData:\n", + " with Reader(asset) as src:\n", + " return src.tile(*args, **kwargs)\n", + "\n", + "\n", + "catalog = Client.open('https://planetarycomputer.microsoft.com/api/stac/v1')\n", + "results = catalog.search(\n", + " limit=10,\n", + " max_items=100,\n", + " bbox=bbox,\n", + " collections=[\"naip\"],\n", + " fields={\"include\":[\"assets.image.href\"], \"exclude\":[\"links\"]}\n", + ")\n", + "print(results.url_with_parameters())\n", + "items=results.items_as_dicts()\n", + "\n", + "assets = (item['assets']['image']['href'] for item in items)\n", + "\n", + "img, used = mosaic_reader(assets, reader, x, y, z, threads=4)\n", + "print(len(used))\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python-3.10.12", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/rio_tiler/mosaic/reader.py b/rio_tiler/mosaic/reader.py index 9729b43b..10fcb83b 100644 --- a/rio_tiler/mosaic/reader.py +++ b/rio_tiler/mosaic/reader.py @@ -2,7 +2,18 @@ import warnings from inspect import isclass -from typing import Any, Callable, List, Optional, Sequence, Tuple, Type, Union, cast +from typing import ( + Any, + Callable, + Iterator, + List, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, +) import numpy from rasterio.crs import CRS @@ -23,7 +34,7 @@ def mosaic_reader( # noqa: C901 - mosaic_assets: Sequence, + mosaic_assets: Union[Iterator, Sequence], reader: Callable[..., ImageData], *args: Any, pixel_selection: Union[Type[MosaicMethodBase], MosaicMethodBase] = FirstMethod, @@ -36,7 +47,7 @@ def mosaic_reader( # noqa: C901 Args: - mosaic_assets (sequence): List of assets. + mosaic_assets (Sequence or Iterator): List of assets. reader (callable): Reader function. The function MUST take `(asset, *args, **kwargs)` as arguments, and MUST return an ImageData. args (Any): Argument to forward to the reader function. pixel_selection (MosaicMethod, optional): Instance of MosaicMethodBase class. Defaults to `rio_tiler.mosaic.methods.defaults.FirstMethod`. @@ -76,9 +87,11 @@ def mosaic_reader( # noqa: C901 "'rio_tiler.mosaic.methods.base.MosaicMethodBase'" ) - if not chunk_size: + if not isinstance(mosaic_assets, Iterator) and not chunk_size: chunk_size = threads if threads > 1 else len(mosaic_assets) + chunk_size = chunk_size or threads + assets_used: List = [] crs: Optional[CRS] bounds: Optional[BBox] diff --git a/rio_tiler/utils.py b/rio_tiler/utils.py index 281a0b58..df1e8ea0 100644 --- a/rio_tiler/utils.py +++ b/rio_tiler/utils.py @@ -1,8 +1,20 @@ """rio_tiler.utils: utility functions.""" +import itertools import warnings from io import BytesIO -from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union +from typing import ( + Any, + Dict, + Generator, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, + Union, +) import numpy import rasterio @@ -26,10 +38,13 @@ from rio_tiler.types import BBox, ColorMapType, IntervalTuple, RIOResampling -def _chunks(my_list: Sequence, chuck_size: int) -> Generator[Sequence, None, None]: +def _chunks(my_list: Iterable, chuck_size: int) -> Generator[Sequence, None, None]: """Yield successive n-sized chunks from l.""" - for i in range(0, len(my_list), chuck_size): - yield my_list[i : i + chuck_size] + if not isinstance(my_list, Iterator): + my_list = iter(my_list) + + while chunk := tuple(itertools.islice(my_list, chuck_size)): + yield chunk def get_array_statistics( diff --git a/tests/test_mosaic.py b/tests/test_mosaic.py index 5fb1a2d3..e0e2ca61 100644 --- a/tests/test_mosaic.py +++ b/tests/test_mosaic.py @@ -286,6 +286,20 @@ class aClass(object): assert m.dtype == "uint8" +def test_mosaic_tiler_iter(): + """Test mosaic tiler with iterator input.""" + assets_iter = iter(assets) + + (t, m), assets_used = mosaic.mosaic_reader(assets_iter, _read_tile, x, y, z) + assert t.shape == (3, 256, 256) + assert m.shape == (256, 256) + assert m.all() + # Should only have value of 1 + assert numpy.unique(t[0, m == 255]).tolist() == [1] + assert t.dtype == "uint16" + assert m.dtype == "uint8" + + def mock_rasterio_open(asset): """Mock rasterio Open.""" assert asset.startswith("http://somewhere-over-the-rainbow.io")