async-geotiff 0.1.0b3__py3-none-any.whl → 0.1.0b5__py3-none-any.whl

async_geotiff/_overview.py

@@ -1,19 +1,25 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from functools import cached_property
 from typing import TYPE_CHECKING
 
 from affine import Affine
 
+from async_geotiff._fetch import FetchTileMixin
+from async_geotiff._read import ReadMixin
+from async_geotiff._transform import TransformMixin
+
 if TYPE_CHECKING:
-    from async_tiff import GeoKeyDirectory, ImageFileDirectory
+    from async_tiff import TIFF, GeoKeyDirectory, ImageFileDirectory
+    from pyproj.crs import CRS
 
     from async_geotiff import GeoTIFF
 
+# ruff: noqa: SLF001
+
 
 @dataclass(init=False, frozen=True, kw_only=True, eq=False, repr=False)
-class Overview:
+class Overview(ReadMixin, FetchTileMixin, TransformMixin):
     """An overview level of a Cloud-Optimized GeoTIFF image."""
 
     _geotiff: GeoTIFF
@@ -24,13 +30,13 @@ class Overview:
     """The GeoKeyDirectory of the primary IFD.
     """
 
-    _ifd: tuple[int, ImageFileDirectory]
+    _ifd: ImageFileDirectory
     """The IFD for this overview level.
 
     (positional index of the IFD in the TIFF file, IFD object)
     """
 
-    _mask_ifd: tuple[int, ImageFileDirectory] | None
+    _mask_ifd: ImageFileDirectory | None
     """The IFD for the mask associated with this overview level, if any.
 
     (positional index of the IFD in the TIFF file, IFD object)
@@ -42,8 +48,8 @@ class Overview:
         *,
         geotiff: GeoTIFF,
         gkd: GeoKeyDirectory,
-        ifd: tuple[int, ImageFileDirectory],
-        mask_ifd: tuple[int, ImageFileDirectory] | None,
+        ifd: ImageFileDirectory,
+        mask_ifd: ImageFileDirectory | None,
     ) -> Overview:
         instance = cls.__new__(cls)
 
@@ -55,13 +61,38 @@ class Overview:
 
         return instance
 
+    @property
+    def _tiff(self) -> TIFF:
+        """A reference to the underlying TIFF object."""
+        return self._geotiff._tiff
+
+    @property
+    def crs(self) -> CRS:
+        """The coordinate reference system of the overview."""
+        return self._geotiff.crs
+
     @property
     def height(self) -> int:
         """The height of the overview in pixels."""
-        return self._ifd[1].image_height
+        return self._ifd.image_height
 
-    @cached_property
-    def transform(self) -> Affine:
+    @property
+    def nodata(self) -> int | float | None:
+        """The nodata value for the overview, if any."""
+        return self._geotiff.nodata
+
+    @property
+    def tile_height(self) -> int:
+        """The height in pixels per tile of the overview."""
+        return self._ifd.tile_height or self.height
+
+    @property
+    def tile_width(self) -> int:
+        """The width in pixels per tile of the overview."""
+        return self._ifd.tile_width or self.width
+
+    @property
+    def transform(self) -> Affine:  # type: ignore[override]
         """The affine transform mapping pixel coordinates to geographic coordinates.
 
         Returns:
@@ -70,9 +101,9 @@ class Overview:
         """
         full_transform = self._geotiff.transform
 
-        overview_width = self._ifd[1].image_width
+        overview_width = self._ifd.image_width
         full_width = self._geotiff.width
-        overview_height = self._ifd[1].image_height
+        overview_height = self._ifd.image_height
         full_height = self._geotiff.height
 
         scale_x = full_width / overview_width
@@ -83,4 +114,4 @@ class Overview:
     @property
     def width(self) -> int:
         """The width of the overview in pixels."""
-        return self._ifd[1].image_width
+        return self._ifd.image_width

async_geotiff/_read.py

@@ -0,0 +1,185 @@
+"""Higher-level read utilities for cross-tile operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+import numpy as np
+from affine import Affine
+
+from async_geotiff._array import Array
+from async_geotiff._fetch import HasTiffReference
+from async_geotiff._windows import Window
+from async_geotiff.exceptions import WindowError
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+    from async_geotiff._tile import Tile
+
+
+class CanFetchTiles(HasTiffReference, Protocol):
+    """Protocol for objects that can fetch tiles."""
+
+    @property
+    def height(self) -> int:
+        """The height of the image in pixels."""
+        ...
+
+    @property
+    def width(self) -> int:
+        """The width of the image in pixels."""
+        ...
+
+    async def fetch_tiles(
+        self,
+        xs: list[int],
+        ys: list[int],
+    ) -> list[Tile]: ...
+
+
+class ReadMixin:
+    async def read(
+        self: CanFetchTiles,
+        *,
+        window: Window | None = None,
+    ) -> Array:
+        """Read pixel data for a window region.
+
+        This method fetches all tiles that intersect the given window and
+        stitches them together, returning only the pixels within the window.
+
+        Args:
+            window: A Window object defining the pixel region to read.
+                If None, the entire image is read.
+
+        Returns:
+            An Array containing the pixel data for the requested window.
+
+        Raises:
+            WindowError: If the window extends outside the image bounds.
+
+        """
+        return await read(self, window=window)
+
+
+async def read(
+    self: CanFetchTiles,
+    *,
+    window: Window | None = None,
+) -> Array:
+    if isinstance(window, Window):
+        win = window
+    else:
+        win = Window(col_off=0, row_off=0, width=self.width, height=self.height)
+
+    # Most validation occurred in construction of Window; here we just check against
+    # image size
+    if win.col_off + win.width > self.width or win.row_off + win.height > self.height:
+        raise WindowError(
+            f"Window extends outside image bounds.\n"
+            f"Window: cols={win.col_off}:{win.col_off + win.width}, "
+            f"rows={win.row_off}:{win.row_off + win.height}.\n"
+            f"Image size: {self.height}x{self.width}",
+        )
+
+    # Calculate which tiles we need to fetch
+    tile_x_start = win.col_off // self.tile_width
+    tile_x_stop = (win.col_off + win.width - 1) // self.tile_width + 1
+    tile_y_start = win.row_off // self.tile_height
+    tile_y_stop = (win.row_off + win.height - 1) // self.tile_height + 1
+
+    # Build list of tile coordinates
+    xs: list[int] = []
+    ys: list[int] = []
+    for tx in range(tile_x_start, tile_x_stop):
+        for ty in range(tile_y_start, tile_y_stop):
+            xs.append(tx)
+            ys.append(ty)
+
+    tiles = await self.fetch_tiles(xs, ys)
+
+    num_bands = tiles[0].array.count
+    dtype = tiles[0].array.data.dtype
+
+    # Create output array and mask array
+    output_data = np.empty((num_bands, win.height, win.width), dtype=dtype)
+    output_mask: NDArray[np.bool_] | None = None
+    if self._mask_ifd is not None:
+        output_mask = np.ones((win.height, win.width), dtype=np.bool_)
+
+    assemble_tiles(
+        tiles=tiles,
+        window=win,
+        tile_width=self.tile_width,
+        tile_height=self.tile_height,
+        output_data=output_data,
+        output_mask=output_mask,
+    )
+
+    window_transform = self.transform * Affine.translation(
+        win.col_off,
+        win.row_off,
+    )
+
+    return Array(
+        data=output_data,
+        mask=output_mask,
+        width=win.width,
+        height=win.height,
+        count=num_bands,
+        transform=window_transform,
+        crs=self.crs,
+    )
+
+
+def assemble_tiles(  # noqa: PLR0913
+    *,
+    tiles: list[Tile],
+    window: Window,
+    tile_width: int,
+    tile_height: int,
+    output_data: NDArray[np.generic],
+    output_mask: NDArray[np.bool_] | None,
+) -> None:
+    """Assemble multiple tiles into output arrays.
+
+    This function copies data from tiles into the appropriate positions
+    in the output arrays, handling partial tiles at window boundaries.
+    """
+    for tile in tiles:
+        # Create a window for this tile's position in image coordinates
+        tile_window = Window(
+            col_off=tile.x * tile_width,
+            row_off=tile.y * tile_height,
+            width=tile.array.width,
+            height=tile.array.height,
+        )
+
+        # Calculate the intersection between tile and target window
+        overlap = window.intersection(tile_window)
+
+        # Calculate source slice within the tile
+        src_col_start = overlap.col_off - tile_window.col_off
+        src_col_stop = src_col_start + overlap.width
+        src_row_start = overlap.row_off - tile_window.row_off
+        src_row_stop = src_row_start + overlap.height
+
+        # Calculate destination slice within the output
+        dst_col_start = overlap.col_off - window.col_off
+        dst_col_stop = dst_col_start + overlap.width
+        dst_row_start = overlap.row_off - window.row_off
+        dst_row_stop = dst_row_start + overlap.height
+
+        # Copy data and mask if present
+        output_data[
+            :,
+            dst_row_start:dst_row_stop,
+            dst_col_start:dst_col_stop,
+        ] = tile.array.data[:, src_row_start:src_row_stop, src_col_start:src_col_stop]
+
+        if output_mask is not None and tile.array.mask is not None:
+            output_mask[
+                dst_row_start:dst_row_stop,
+                dst_col_start:dst_col_stop,
+            ] = tile.array.mask[src_row_start:src_row_stop, src_col_start:src_col_stop]

async_geotiff/_tile.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from async_tiff import ImageFileDirectory
+
+    from async_geotiff._array import Array
+
+
+@dataclass(frozen=True, kw_only=True, eq=False)
+class Tile:
+    """A tile from a GeoTIFF, containing array data and grid position."""
+
+    x: int
+    """The tile column index in the GeoTIFF or overview."""
+
+    y: int
+    """The tile row index in the GeoTIFF or overview."""
+
+    _ifd: ImageFileDirectory
+    """A reference to the IFD this tile belongs to."""
+
+    array: Array
+    """The array data for this tile."""

async_geotiff/_transform.py

@@ -0,0 +1,90 @@
+"""Mixin class for coordinate transformation methods."""
+
+from __future__ import annotations
+
+from math import floor
+from typing import TYPE_CHECKING, Literal, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from affine import Affine
+
+
+class HasTransform(Protocol):
+    """Protocol for objects that have an affine transform."""
+
+    @property
+    def transform(self) -> Affine: ...
+
+
+class TransformMixin:
+    """Mixin providing coordinate transformation methods.
+
+    Classes using this mixin must implement HasTransform.
+    """
+
+    def index(
+        self: HasTransform,
+        x: float,
+        y: float,
+        op: Callable[[float], int] = floor,
+    ) -> tuple[int, int]:
+        """Get the (row, col) index of the pixel containing (x, y).
+
+        Args:
+            x: x value in coordinate reference system.
+            y: y value in coordinate reference system.
+            op: Function to convert fractional pixels to whole numbers
+                (floor, ceiling, round). Defaults to math.floor.
+
+        Returns:
+            (row index, col index)
+
+        """
+        inv_transform = ~self.transform
+        # Affine * (x, y) returns tuple[float, float] for 2D coordinates
+        col_frac, row_frac = inv_transform * (x, y)  # type: ignore[misc]
+
+        return (op(row_frac), op(col_frac))
+
+    def xy(
+        self: HasTransform,
+        row: int,
+        col: int,
+        offset: Literal["center", "ul", "ur", "ll", "lr"] = "center",
+    ) -> tuple[float, float]:
+        """Get the coordinates (x, y) of a pixel at (row, col).
+
+        The pixel's center is returned by default, but a corner can be returned
+        by setting `offset` to one of `"ul"`, `"ur"`, `"ll"`, `"lr"`.
+
+        Args:
+            row: Pixel row.
+            col: Pixel column.
+            offset: Determines if the returned coordinates are for the center of the
+                pixel or for a corner.
+
+        Returns:
+            (x, y) coordinates in the dataset's CRS.
+
+        """
+        if offset == "center":
+            c = col + 0.5
+            r = row + 0.5
+        elif offset == "ul":
+            c = col
+            r = row
+        elif offset == "ur":
+            c = col + 1
+            r = row
+        elif offset == "ll":
+            c = col
+            r = row + 1
+        elif offset == "lr":
+            c = col + 1
+            r = row + 1
+        else:
+            raise ValueError(f"Invalid offset value: {offset}")
+
+        return self.transform * (c, r)

async_geotiff/_windows.py

@@ -0,0 +1,76 @@
+"""Window utilities for defining rectangular subsets of rasters."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from async_geotiff.exceptions import WindowError
+
+
+@dataclass(frozen=True, slots=True)
+class Window:
+    """A rectangular subset of a raster.
+
+    Windows define pixel regions using column/row offsets and dimensions.
+    This class is similar to rasterio's [Window][rasterio.windows.Window] but supports
+    integer offsets and ranges only.
+    """
+
+    col_off: int
+    """The column offset (x position of the left edge)."""
+
+    row_off: int
+    """The row offset (y position of the top edge)."""
+
+    width: int
+    """The width in pixels (number of columns)."""
+
+    height: int
+    """The height in pixels (number of rows)."""
+
+    def __post_init__(self) -> None:
+        """Validate window dimensions."""
+        if self.col_off < 0 or self.row_off < 0:
+            raise WindowError(
+                f"Window start indices must be non-negative, "
+                f"got col_off={self.col_off}, row_off={self.row_off}",
+            )
+
+        if self.width <= 0:
+            raise WindowError(f"Window width must be positive, got {self.width}")
+
+        if self.height <= 0:
+            raise WindowError(f"Window height must be positive, got {self.height}")
+
+    def __repr__(self) -> str:
+        """Return a nicely formatted representation string."""
+        return (
+            f"async_geotiff.Window(col_off={self.col_off}, row_off={self.row_off}, "
+            f"width={self.width}, height={self.height})"
+        )
+
+    def intersection(self, other: Window) -> Window:
+        """Compute the intersection with another window.
+
+        Args:
+            other: Another Window object.
+
+        Returns:
+            A new Window representing the overlapping region.
+
+        Raises:
+            WindowError: If windows do not intersect.
+
+        """
+        col_off = max(self.col_off, other.col_off)
+        row_off = max(self.row_off, other.row_off)
+        col_stop = min(self.col_off + self.width, other.col_off + other.width)
+        row_stop = min(self.row_off + self.height, other.row_off + other.height)
+
+        width = col_stop - col_off
+        height = row_stop - row_off
+
+        if width <= 0 or height <= 0:
+            raise WindowError(f"Windows do not intersect: {self} and {other}")
+
+        return Window(col_off=col_off, row_off=row_off, width=width, height=height)

async_geotiff/colormap.py

@@ -0,0 +1,104 @@
+"""High-level Colormap class for GeoTIFF colormaps."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from async_tiff import Colormap as AsyncTiffColormap
+    from numpy.typing import NDArray
+
+
+@dataclass(frozen=True, kw_only=True, eq=False)
+class Colormap:
+    """A representation of a GeoTIFF colormap.
+
+    GeoTIFF colormaps
+    """
+
+    _cmap: AsyncTiffColormap
+    """The colormap data held in Rust, accessible via the buffer protocol.
+
+    Has shape `(N, 3)` and is of data type uint16.
+    """
+
+    _nodata: int | float | None
+    """The nodata value from gdal_nodata, if set."""
+
+    def as_array(self, *, dtype: type[np.uint8 | np.uint16] = np.uint8) -> NDArray:
+        """Return the colormap as a NumPy array with shape (N, 3) and dtype uint16.
+
+        Each row corresponds to a color entry in the colormap, with columns
+        representing the Red, Green, and Blue components respectively.
+
+        This is the most efficient way to access and apply the colormap data.
+
+        ```py
+        geotiff = await GeoTIFF.open(...)
+        array = await geotiff.fetch_tile(0, 0)
+
+        colormap = geotiff.colormap
+        colormap_array = colormap.as_array()
+
+        rgb_data = colormap_array[array.data[0]]
+        # A 3D array with shape (height, width, 3)
+        ```
+
+        Returns:
+            A NumPy array representation of the colormap.
+
+        """
+        cmap_array = np.asarray(self._cmap)
+        if dtype == np.uint8:
+            return (cmap_array >> 8).astype(np.uint8)
+        if dtype == np.uint16:
+            return cmap_array
+        raise ValueError("dtype must be either np.uint8 or np.uint16.")
+
+    def as_dict(
+        self,
+        *,
+        dtype: type[np.uint8 | np.uint16] = np.uint8,
+    ) -> dict[int, tuple[int, int, int]]:
+        """Return the colormap as a dictionary mapping indices to RGB tuples.
+
+        Returns:
+            A dictionary where keys are indices and values are tuples of
+                (Red, Green, Blue) components.
+
+        """
+        cmap_array = self.as_array(dtype=dtype)
+        return {
+            int(idx): (int(r), int(g), int(b))
+            for idx, (r, g, b) in enumerate(cmap_array)
+        }
+
+    def as_rasterio(self) -> dict[int, tuple[int, int, int, int]]:
+        """Return the colormap as a mapping to 8-bit RGBA colors.
+
+        This returns a colormap in the same format as rasterio's
+        [`DatasetReader.colormap`][rasterio.io.DatasetReader.colormap] method.
+
+        This is the same as
+        [`Colormap.as_dict`][async_geotiff.colormap.Colormap.as_dict] with:
+
+        - `dtype` set to `np.uint8`
+        - an added alpha channel set to 255, **except** for the nodata value, if
+          defined, which has an alpha of 0.
+
+        Returns:
+            Mapping of color index value (starting at 0) to RGBA color as a 4-element
+                tuple.
+
+        """
+        cmap_array = self.as_array(dtype=np.uint8)
+        cmap_dict: dict[int, tuple[int, int, int, int]] = {}
+
+        for idx, (r, g, b) in enumerate(cmap_array):
+            alpha = 255 if self._nodata is None or idx != self._nodata else 0
+            cmap_dict[int(idx)] = (int(r), int(g), int(b), alpha)
+
+        return cmap_dict

async_geotiff/exceptions.py

@@ -0,0 +1,5 @@
+"""Exceptions for async_geotiff package."""
+
+
+class WindowError(Exception):
+    """Exception raised for window-related errors."""

async_geotiff/tms.py

@@ -63,8 +63,8 @@ def generate_tms(
 
     for idx, overview in enumerate(reversed(geotiff.overviews)):
         overview_tr = overview.transform
-        blockxsize = overview._ifd[1].tile_width  # noqa: SLF001
-        blockysize = overview._ifd[1].tile_height  # noqa: SLF001
+        blockxsize = overview._ifd.tile_width  # noqa: SLF001
+        blockysize = overview._ifd.tile_height  # noqa: SLF001
 
         if blockxsize is None or blockysize is None:
             raise ValueError("GeoTIFF overviews must be tiled to generate a TMS.")