diff --git a/geoutils/raster/raster.py b/geoutils/raster/raster.py index aaedfbf0..eee64126 100644 --- a/geoutils/raster/raster.py +++ b/geoutils/raster/raster.py @@ -171,6 +171,7 @@ def _load_rio( transform: Affine | None = None, shape: tuple[int, int] | None = None, out_count: int | None = None, + roi: dict[str, int] | None = None, **kwargs: Any, ) -> MArrayNum: r""" @@ -185,6 +186,7 @@ def _load_rio( :param transform: Create a window from the given transform (to read only parts of the raster) :param shape: Expected shape of the read ndarray. Must be given together with the `transform` argument. :param out_count: Specify the count for a downsampled version (to be used with kwargs out_shape). + :param roi: Optional pixel-based region of interest (dict with keys: left, bottom, right, top). :raises ValueError: If only one of ``transform`` and ``shape`` are given. @@ -197,9 +199,19 @@ def _load_rio( * window : to load a cropped version * resampling : to set the resampling algorithm """ + window = None + + # If a roi is passed, set up the corresponding window + if roi is not None: + window = rio.windows.Window( + col_off=roi["left"], + row_off=dataset.height - roi["top"], + width=roi["right"] - roi["left"], + height=roi["top"] - roi["bottom"], + ) + # If out_shape is passed, no need to account for transform and shape - if kwargs["out_shape"] is not None: - window = None + elif kwargs.get("out_shape") is not None: # If multi-band raster, the out_shape needs to contain the count if out_count is not None and out_count > 1: kwargs["out_shape"] = (out_count, *kwargs["out_shape"]) @@ -213,8 +225,6 @@ def _load_rio( window = rio.windows.Window(col_off, row_off, *shape[::-1]) elif sum(param is None for param in [shape, transform]) == 1: raise ValueError("If 'shape' or 'transform' is provided, BOTH must be given.") - else: - window = None if indexes is None: if only_mask: @@ -356,6 +366,7 @@ def __init__( silent: bool = True, downsample: Number = 1, nodata: int | float | None = None, + roi: dict[str, int] | None = None, ) -> None: """ Instantiate a raster from a filename or rasterio dataset. @@ -367,6 +378,7 @@ def __init__( :param silent: Whether to parse metadata silently or with console output. :param downsample: Downsample the array once loaded by a round factor. Default is no downsampling. :param nodata: Nodata value to be used (overwrites the metadata). Default reads from metadata. + :param roi: Optional pixel-based region of interest (dict with keys: left, bottom, right, top). """ self._driver: str | None = None self._name: str | None = None @@ -377,6 +389,7 @@ def __init__( self._transform: affine.Affine | None = None self._crs: CRS | None = None self._nodata: int | float | None = nodata + self._roi: dict[str, int] | None = roi self._bands = bands self._bands_loaded: int | tuple[int, ...] | None = None self._masked = True @@ -391,6 +404,14 @@ def __init__( self._downsample: int | float = 1 self._area_or_point: Literal["Area", "Point"] | None = None + # Check that roi contains the needed keys + if roi is not None: + required_roi_keys = {"left", "bottom", "right", "top"} + if not required_roi_keys.issubset(roi.keys()): + raise ValueError( + "The roi parameter must contain the following keys : 'left', 'right', 'top' and 'bottom'" + ) + # This is for Raster.from_array to work. if isinstance(filename_or_dataset, dict): @@ -412,7 +433,7 @@ def __init__( self.crs: rio.crs.CRS = filename_or_dataset["crs"] for key in filename_or_dataset: - if key in ["data", "transform", "crs", "nodata", "area_or_point", "tags"]: + if key in ["data", "transform", "crs", "nodata", "area_or_point", "tags", "roi"]: continue setattr(self, key, filename_or_dataset[key]) return @@ -421,7 +442,17 @@ def __init__( if isinstance(filename_or_dataset, Raster): for key in filename_or_dataset.__dict__: setattr(self, key, filename_or_dataset.__dict__[key]) - return + + # if a roi is passed, crop the raster + if roi is not None: + crop_extent = [ + self.transform.c + roi["left"] * self.transform.a, # xmin + self.transform.f + (self.height - roi["bottom"]) * self.transform.e, # ymin + self.transform.c + roi["right"] * self.transform.a, # xmax + self.transform.f + (self.height - roi["top"]) * self.transform.e, # ymax + ] + self.crop(crop_extent, inplace=True) + # Image is a file on disk. elif isinstance(filename_or_dataset, (str, pathlib.Path, rio.io.DatasetReader, rio.io.MemoryFile)): # ExitStack is used instead of "with rio.open(filename_or_dataset) as ds:". @@ -472,7 +503,16 @@ def __init__( # Downsampled image size if not isinstance(downsample, (int, float)): raise TypeError("downsample must be of type int or float.") - if downsample == 1: + if self._roi is not None: + new_transform = rio.transform.from_origin( + self.transform.c + self._roi["left"] * self.transform.a, + self.transform.f + (self.height - self._roi["top"]) * self.transform.e, + self.transform.a, + -self.transform.e, + ) + self.transform = new_transform + out_shape = (self._roi["top"] - self._roi["bottom"], self._roi["right"] - self._roi["left"]) + elif downsample == 1: out_shape = (self.height, self.width) else: down_width = int(np.ceil(self.width / downsample)) @@ -495,6 +535,7 @@ def __init__( masked=self._masked, out_shape=out_shape, out_count=count, + roi=self._roi, ) # type: ignore # Probably don't want to use set_nodata that can update array, setting self._nodata is sufficient @@ -791,6 +832,7 @@ def _load_only_mask(self, bands: int | list[int] | None = None, **kwargs: Any) - shape=self.shape, out_shape=self._out_shape, out_count=out_count, + roi=self._roi, **kwargs, ) @@ -847,6 +889,7 @@ def load(self, bands: int | list[int] | None = None, **kwargs: Any) -> None: shape=self.shape, out_shape=self._out_shape, out_count=self._out_count, + roi=self._roi, **kwargs, ) @@ -869,6 +912,7 @@ def from_array( area_or_point: Literal["Area", "Point"] | None = None, tags: dict[str, Any] = None, cast_nodata: bool = True, + roi: dict[str, int] = None, ) -> RasterType: """Create a raster from a numpy array and the georeferencing information. @@ -883,6 +927,7 @@ def from_array( :param tags: Metadata stored in a dictionary. :param cast_nodata: Automatically cast nodata value to the default nodata for the new array type if not compatible. If False, will raise an error when incompatible. + :param roi: Optional dictionary specifying the region of interest (ROI) with keys: left, bottom, right, top. :returns: Raster created from the provided array and georeferencing. @@ -903,6 +948,17 @@ def from_array( if cast_nodata: nodata = _cast_nodata(data.dtype, nodata) + if roi is not None: + height = data.shape[-2] + new_transform = rio.transform.from_origin( + transform[2] + roi["left"] * transform[0], + transform[5] + (height - roi["top"]) * transform[4], + transform[0], + -transform[4], + ) + transform = new_transform + data = data[..., height - roi["top"] : height - roi["bottom"], roi["left"] : roi["right"]] + # If the data was transformed into boolean, re-initialize as a Mask subclass # Typing: we can specify this behaviour in @overload once we add the NumPy plugin of MyPy if data.dtype == bool: