Repository: ungarj/mapchete
Pull request #513 (2022-11-28 12:18:08)

convert ReferencedRaster to class to allow adding methods


Patch:
diff --git a/mapchete/io/_geometry_operations.py b/mapchete/io/_geometry_operations.py
index 1becd64..0fdc1d0 100644
--- a/mapchete/io/_geometry_operations.py
+++ b/mapchete/io/_geometry_operations.py
@@ -78,7 +78,11 @@ def reproject_geometry(
     """
     src_crs = validate_crs(src_crs)
     dst_crs = validate_crs(dst_crs)
-    if not isinstance(geometry, base.BaseGeometry):  # pragma: no cover
+    try:
+        geometry = (
+            geometry if isinstance(geometry, base.BaseGeometry) else shape(geometry)
+        )
+    except Exception:  # pragma: no cover
         raise TypeError(f"invalid geometry type: {type(geometry)}")
 
     def _reproject_geom(geometry, src_crs, dst_crs):
diff --git a/mapchete/io/raster.py b/mapchete/io/raster.py
index c0e1b2f..55994eb 100644
--- a/mapchete/io/raster.py
+++ b/mapchete/io/raster.py
@@ -1,7 +1,6 @@
 """Wrapper functions around rasterio and useful raster functions."""
 
 from affine import Affine
-from collections import namedtuple
 import itertools
 import logging
 import numpy as np
@@ -14,11 +13,13 @@ from rasterio.errors import RasterioIOError
 from rasterio.io import MemoryFile
 from rasterio.transform import from_bounds as affine_from_bounds
 from rasterio.vrt import WarpedVRT
-from rasterio.warp import reproject, calculate_default_transform
+from rasterio.warp import reproject
 from rasterio.windows import from_bounds
+from shapely.geometry import box, mapping
 from tempfile import NamedTemporaryFile
 from tilematrix import clip_geometry_to_srs_bounds, Shape, Bounds
 from types import GeneratorType
+from typing import List, Tuple, Union
 import warnings
 
 from mapchete.errors import MapcheteIOError
@@ -38,7 +39,113 @@ from mapchete.validate import validate_write_window_params
 
 logger = logging.getLogger(__name__)
 
-ReferencedRaster = namedtuple("ReferencedRaster", ("data", "affine", "bounds", "crs"))
+
+class ReferencedRaster:
+    """
+    A loose in-memory representation of a rasterio dataset.
+
+    Useful to ship cached raster files between worker nodes.
+    """
+
+    def __init__(
+        self,
+        data: Union[np.ndarray, ma.masked_array] = None,
+        transform: Affine = None,
+        bounds: Union[List[float], Tuple[float], Bounds] = None,
+        crs: Union[str, rasterio.crs.CRS] = None,
+        nodata: Union[int, float] = None,
+        driver: str = None,
+        **kwargs,
+    ):
+        if data.ndim == 2:
+            self.count = 1
+            self.height, self.width = data.shape
+        elif data.ndim == 3:
+            self.count, self.height, self.width = data.shape
+        else:  # pragma: no cover
+            raise TypeError("input array must be 2- or 3-dimensional")
+        transform = transform or kwargs.get("affine")
+        if transform is None:  # pragma: no cover
+            raise ValueError("georeference given")
+        self.data = data
+        self.driver = driver
+        self.dtype = self.data.dtype
+        self.nodata = nodata
+        self.crs = crs
+        self.transform = self.affine = transform
+        self.bounds = bounds
+        self.__geo_interface__ = mapping(box(*self.bounds))
+
+    @property
+    def meta(self) -> dict:
+        return {
+            "driver": self.driver,
+            "dtype": self.dtype,
+            "nodata": self.nodata,
+            "width": self.width,
+            "height": self.height,
+            "count": self.count,
+            "crs": self.crs,
+            "transform": self.transform,
+        }
+
+    def read(
+        self,
+        indexes: Union[int, List[int]] = None,
+        tile: BufferedTile = None,
+        resampling: str = "nearest",
+    ) -> np.ndarray:
+        """Either read full array or resampled to tile."""
+        # select bands using band indexes
+        if indexes is None or self.data.ndim == 2:
+            band_selection = self.data
+        else:
+            band_selection = self._stack(
+                [self.data[i - 1] for i in self._get_band_indexes(indexes)]
+            )
+
+        # return either full array or a window resampled to tile
+        if tile is None:
+            return band_selection
+        else:
+            return resample_from_array(
+                in_raster=band_selection,
+                in_affine=self.transform,
+                in_crs=self.crs,
+                nodataval=self.nodata,
+                nodata=self.nodata,
+                out_tile=tile,
+                resampling=resampling,
+            )
+
+    def _get_band_indexes(self, indexes: Union[List[int], int] = None) -> List[int]:
+        """Return valid band indexes."""
+        if isinstance(indexes, int):
+            return [indexes]
+        else:
+            return indexes
+
+    def _stack(self, *args) -> np.ndarray:
+        """return stack of numpy or numpy.masked depending on array type"""
+        return (
+            ma.stack(*args)
+            if isinstance(self.data, ma.masked_array)
+            else np.stack(*args)
+        )
+
+    @staticmethod
+    def from_rasterio(src, masked: bool = True) -> "ReferencedRaster":
+        return ReferencedRaster(
+            data=src.read(masked=masked).copy(),
+            transform=src.transform,
+            bounds=src.bounds,
+            crs=src.crs,
+        )
+
+    @staticmethod
+    def from_file(path, masked: bool = True) -> "ReferencedRaster":
+        with rasterio.open(path) as src:
+            return ReferencedRaster.from_rasterio(src, masked=masked)
 
 
 def read_raster_window(


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
index 76a8842..114a804 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -4,17 +4,20 @@ from collections import namedtuple
 import datetime
 import os
 import pytest
+import rasterio
 from shapely import wkt
+from shapely.geometry import box
 import shutil
 import uuid
-import yaml
 
 from tilematrix import Bounds, GridDefinition
 
 from mapchete.cli.default.serve import create_app
 from mapchete._executor import DaskExecutor
 from mapchete.io import fs_from_path
-from mapchete.testing import ProcessFixture, dict_from_mapchete
+from mapchete.io.vector import reproject_geometry
+from mapchete.testing import ProcessFixture
+from mapchete.tile import BufferedTilePyramid
 
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
@@ -105,6 +108,16 @@ def s2_band():
     return os.path.join(TESTDATA_DIR, "s2_band.tif")
 
 
+@pytest.fixture
+def s2_band_tile():
+    tp = BufferedTilePyramid("geodetic")
+    with rasterio.open(os.path.join(TESTDATA_DIR, "s2_band.tif")) as src:
+        rr_center = reproject_geometry(
+            geometry=box(*src.bounds), src_crs=src.crs, dst_crs=tp.crs
+        ).centroid
+        return next(tp.tiles_from_geom(rr_center, 13))
+
+
 @pytest.fixture
 def s2_band_jp2():
     """
diff --git a/test/test_io_raster.py b/test/test_io_raster.py
index 1375f67..97cb9bd 100644
--- a/test/test_io_raster.py
+++ b/test/test_io_raster.py
@@ -29,6 +29,7 @@ from mapchete.io.raster import (
     convert_raster,
 )
 from mapchete.tile import BufferedTilePyramid
+from mapchete.io.vector import reproject_geometry
 
 
 def test_read_raster_window(dummy1_tif, minmax_zoom):
@@ -835,3 +836,31 @@ def test_convert_raster_other_format_overwrite(cleantopo_br_tif, tmpdir):
     convert_raster(cleantopo_br_tif, out, driver="JP2OpenJPEG", overwrite=True)
     with rasterio.open(out) as src:
         assert not src.read(masked=True).mask.all()
+
+
+def test_referencedraster_meta(s2_band):
+    rr = ReferencedRaster.from_file(s2_band)
+    meta = rr.meta
+    for k in [
+        "driver",
+        "dtype",
+        "nodata",
+        "width",
+        "height",
+        "count",
+        "crs",
+        "transform",
+    ]:
+        assert k in meta
+
+
+@pytest.mark.parametrize("indexes", [None, 1, [1]])
+def test_referencedraster_read_band(s2_band, indexes):
+    rr = ReferencedRaster.from_file(s2_band)
+    assert rr.read(indexes).any()
+
+
+@pytest.mark.parametrize("indexes", [None, 1, [1]])
+def test_referencedraster_read_tile_band(s2_band, indexes, s2_band_tile):
+    rr = ReferencedRaster.from_file(s2_band)
+    assert rr.read(indexes, tile=s2_band_tile).any()