diff --git a/CHANGES.md b/CHANGES.md index 6ad8a0da..1a1f8843 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -10,6 +10,7 @@ * Moved `_dst_geom_in_tms_crs` from Reader to `SpatialMixin` class **breaking change** * Removed use of rasterio's `is_tiled` method * Enable **Alternate** asset's HREF for STAC by using `RIO_TILER_STAC_ALTERNATE_KEY` environment variable +* Adding support for GDAL VRT Connection string for STAC Assets * Improve type hint definition * make `ImageData.rescale` and `ImageData.apply_color_formula` to return `self` diff --git a/rio_tiler/io/stac.py b/rio_tiler/io/stac.py index 6c013b3a..0a05111e 100644 --- a/rio_tiler/io/stac.py +++ b/rio_tiler/io/stac.py @@ -293,6 +293,23 @@ def _get_reader(self, asset_info: AssetInfo) -> Tuple[Type[BaseReader], Dict]: """Get Asset Reader.""" return self.reader, {} + def _parse_vrt_asset(self, asset: str) -> Tuple[str, Optional[str]]: + if asset.startswith("vrt://") and asset not in self.assets: + parsed = urlparse(asset) + if not parsed.netloc: + raise InvalidAssetName( + f"'{asset}' is not valid, couldn't find valid asset" + ) + + if parsed.netloc not in self.assets: + raise InvalidAssetName( + f"'{parsed.netloc}' is not valid, should be one of {self.assets}" + ) + + return parsed.netloc, parsed.query + + return asset, None + def _get_asset_info(self, asset: str) -> AssetInfo: """Validate asset names and return asset's info. @@ -303,6 +320,7 @@ def _get_asset_info(self, asset: str) -> AssetInfo: AssetInfo: STAC asset info. """ + asset, vrt_options = self._parse_vrt_asset(asset) if asset not in self.assets: raise InvalidAssetName( f"'{asset}' is not valid, should be one of {self.assets}" @@ -313,7 +331,7 @@ def _get_asset_info(self, asset: str) -> AssetInfo: info = AssetInfo( url=asset_info.get_absolute_href() or asset_info.href, - metadata=extras, + metadata=extras if not vrt_options else None, ) if STAC_ALTERNATE_KEY and extras.get("alternate"): @@ -328,7 +346,8 @@ def _get_asset_info(self, asset: str) -> AssetInfo: info["env"] = {"GDAL_INGESTED_BYTES_AT_OPEN": head} # https://github.com/stac-extensions/raster - if bands := extras.get("raster:bands"): + if extras.get("raster:bands") and not vrt_options: + bands = extras.get("raster:bands") stats = [ (b["statistics"]["minimum"], b["statistics"]["maximum"]) for b in bands @@ -346,4 +365,7 @@ def _get_asset_info(self, asset: str) -> AssetInfo: "Some statistics data in STAC are invalid, they will be ignored." ) + if vrt_options: + info["url"] = f"vrt://{info['url']}?{vrt_options}" + return info diff --git a/tests/fixtures/gfs.t06z.pgrb2.10p0.f010.grib2 b/tests/fixtures/gfs.t06z.pgrb2.10p0.f010.grib2 new file mode 100644 index 00000000..af97027f Binary files /dev/null and b/tests/fixtures/gfs.t06z.pgrb2.10p0.f010.grib2 differ diff --git a/tests/fixtures/stac_grib.json b/tests/fixtures/stac_grib.json new file mode 100644 index 00000000..09b5096d --- /dev/null +++ b/tests/fixtures/stac_grib.json @@ -0,0 +1,304 @@ +{ + "type": "Feature", + "stac_version": "1.0.0", + "id": "gribfile", + "properties": { + "proj:geometry": { + "type": "Polygon", + "coordinates": [ + [ + [ + -180.125, + -89.875 + ], + [ + 179.875, + -89.875 + ], + [ + 179.875, + 90.125 + ], + [ + -180.125, + 90.125 + ], + [ + -180.125, + -89.875 + ] + ] + ] + }, + "proj:bbox": [ + -180.125, + -89.875, + 179.875, + 90.125 + ], + "proj:shape": [ + 18, + 36 + ], + "proj:transform": [ + 10.0, + 0.0, + -180.125, + 0.0, + -10.0, + 90.125, + 0.0, + 0.0, + 1.0 + ], + "proj:wkt2": "GEOGCS[\"Coordinate System imported from GRIB file\",DATUM[\"unnamed\",SPHEROID[\"Sphere\",6371229,0]],PRIMEM[\"Greenwich\",0],UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],AXIS[\"Latitude\",NORTH],AXIS[\"Longitude\",EAST]]", + "datetime": "2024-09-13T11:08:36.893626Z" + }, + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [ + -180.125, + -89.875 + ], + [ + 179.875, + -89.875 + ], + [ + 179.875, + 90.125 + ], + [ + -180.125, + 90.125 + ], + [ + -180.125, + -89.875 + ] + ] + ] + }, + "links": [], + "assets": { + "asset": { + "href": "gfs.t06z.pgrb2.10p0.f010.grib2", + "raster:bands": [ + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": -12.643815520808767, + "minimum": -20.000003814697266, + "maximum": 32.659996032714844, + "stddev": 12.448285782027321, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": -20.000003814697266, + "max": 32.659996032714844, + "buckets": [ + 449, + 27, + 29, + 31, + 35, + 29, + 15, + 21, + 9, + 3 + ] + } + }, + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": -12.609587132580073, + "minimum": -20.000003814697266, + "maximum": 32.66999435424805, + "stddev": 12.450450829609293, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": -20.000003814697266, + "max": 32.66999435424805, + "buckets": [ + 447, + 27, + 31, + 29, + 39, + 28, + 14, + 21, + 9, + 3 + ] + } + }, + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": -10.566084066483503, + "minimum": -20.000003814697266, + "maximum": 32.76999282836914, + "stddev": 13.981093035778388, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": -20.000003814697266, + "max": 32.76999282836914, + "buckets": [ + 412, + 25, + 31, + 27, + 34, + 49, + 32, + 20, + 11, + 7 + ] + } + }, + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": 20302.916205571022, + "minimum": 24.859272003173828, + "maximum": 24134.859375, + "stddev": 7731.471190491879, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": 24.859272003173828, + "max": 24134.859375, + "buckets": [ + 49, + 24, + 17, + 11, + 8, + 6, + 11, + 11, + 11, + 500 + ] + } + }, + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": -0.13451403068101184, + "minimum": -18.402570724487305, + "maximum": 29.097431182861328, + "stddev": 7.346453975073525, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": -18.402570724487305, + "max": 29.097431182861328, + "buckets": [ + 16, + 47, + 120, + 191, + 153, + 64, + 32, + 15, + 8, + 2 + ] + } + }, + { + "data_type": "float64", + "scale": 1.0, + "offset": 0.0, + "statistics": { + "mean": 0.3850667698676755, + "minimum": -27.066476821899414, + "maximum": 20.933523178100586, + "stddev": 6.63843088754233, + "valid_percent": 0.15432098765432098 + }, + "histogram": { + "count": 11, + "min": -27.066476821899414, + "max": 20.933523178100586, + "buckets": [ + 2, + 7, + 15, + 34, + 106, + 212, + 175, + 64, + 25, + 8 + ] + } + } + ], + "eo:bands": [ + { + "name": "b1", + "description": "1[-] HYBL=\"Hybrid level\"" + }, + { + "name": "b2", + "description": "2[-] HYBL=\"Hybrid level\"" + }, + { + "name": "b3", + "description": "0[-] EATM=\"Entire Atmosphere\"" + }, + { + "name": "b4", + "description": "0[-] SFC=\"Ground or water surface\"" + }, + { + "name": "b5", + "description": "0[-] RESERVED(220) (Reserved for local use)" + }, + { + "name": "b6", + "description": "0[-] RESERVED(220) (Reserved for local use)" + } + ], + "roles": [] + } + }, + "bbox": [ + -180.125, + -89.875, + 179.875, + 90.125 + ], + "stac_extensions": [ + "https://stac-extensions.github.io/projection/v1.1.0/schema.json", + "https://stac-extensions.github.io/raster/v1.1.0/schema.json", + "https://stac-extensions.github.io/eo/v1.1.0/schema.json" + ] +} diff --git a/tests/test_io_stac.py b/tests/test_io_stac.py index 97ed00d0..02f65bd0 100644 --- a/tests/test_io_stac.py +++ b/tests/test_io_stac.py @@ -23,6 +23,7 @@ TileOutsideBounds, ) from rio_tiler.io import BaseReader, Reader, STACReader, XarrayReader +from rio_tiler.io.stac import DEFAULT_VALID_TYPE from rio_tiler.models import BandStatistics from rio_tiler.types import AssetInfo @@ -34,6 +35,7 @@ STAC_RASTER_PATH = os.path.join(PREFIX, "stac_raster.json") STAC_WRONGSTATS_PATH = os.path.join(PREFIX, "stac_wrong_stats.json") STAC_ALTERNATE_PATH = os.path.join(PREFIX, "stac_alternate.json") +STAC_GRIB_PATH = os.path.join(PREFIX, "stac_grib.json") with open(STAC_PATH) as f: item = json.loads(f.read()) @@ -1024,3 +1026,43 @@ def test_alternate_assets(): assert stac._get_asset_info("red")["url"].startswith("s3://") # fall back to href when alternate doesn't exist assert stac._get_asset_info("blue")["url"].startswith("http://") + + +def test_vrt_string_assets(): + """Should work with VRT connection string""" + VALID_TYPE = { + *DEFAULT_VALID_TYPE, + "application/wmo-GRIB2", + } + + with STACReader(STAC_GRIB_PATH, include_asset_types=VALID_TYPE) as stac: + assert stac.assets == ["asset"] + info = stac._get_asset_info("asset") + assert info["url"] + + info_vrt = stac._get_asset_info("vrt://asset") + # without any option there is no need to use the vrt prefix + assert info["url"] == info_vrt["url"] + + info_vrt = stac._get_asset_info("vrt://asset?bands=1") + assert not info["url"] == info_vrt["url"] + assert info_vrt["url"].startswith("vrt://") and info_vrt["url"].endswith( + "?bands=1" + ) + + with pytest.raises(InvalidAssetName): + stac._get_asset_info("vrt://somthing?bands=1") + + with pytest.raises(InvalidAssetName): + stac._get_asset_info("vrt://?bands=1") + + info = stac.info(assets="vrt://asset?bands=1") + assert info["vrt://asset?bands=1"] + assert len(info["vrt://asset?bands=1"].band_metadata) == 1 + + info = stac.info(assets="vrt://asset?bands=1,2") + assert info["vrt://asset?bands=1,2"] + assert len(info["vrt://asset?bands=1,2"].band_metadata) == 2 + + img = stac.preview(assets="vrt://asset?bands=1") + assert img.count == 1