feat: Support loading transparency masks (#36)

* feat: Support loading transparency masks

* improve docstring

Author: kylebarron

fixtures/geotiff-test-data

@@ -22,7 +22,7 @@ classifiers = [
 keywords = ["geotiff", "tiff", "async", "cog", "raster", "gis"]
 dependencies = [
     "affine>=2.4.0",
-    "async-tiff>=0.4.0",
+    "async-tiff>=0.5.0-beta.1",
     "numpy>=2.0",
     "pyproj>=3.3.0",
 ]

pyproject.toml

@@ -22,8 +22,11 @@ class Array(TransformMixin):
     data: NDArray
     """The array data with shape (bands, height, width)."""
 
-    mask: NDArray | None
-    """The mask array with shape (height, width), if any."""
+    mask: NDArray[np.bool_] | None
+    """The mask array with shape (height, width), if any.
+
+    Values of True indicate valid data; False indicates no data.
+    """
 
     width: int
     """The width of the array in pixels."""
@@ -66,8 +69,15 @@ def _create(
             An Array with data in (bands, height, width) order.
 
         """
-        data_arr = np.asarray(data)
-        mask_arr = np.asarray(mask) if mask is not None else None
+        data_arr = np.asarray(data, copy=False)
+        if mask is not None:
+            mask_arr = np.asarray(mask, copy=False).astype(np.bool_, copy=False)
+            assert mask_arr.ndim == 3  # noqa: S101, PLR2004
+            assert mask_arr.shape[2] == 1  # noqa: S101
+            # This assumes it's always (height, width, 1)
+            mask_arr = np.squeeze(mask_arr, axis=2)
+        else:
+            mask_arr = None
 
         assert data_arr.ndim == 3, f"Expected 3D array, got {data_arr.ndim}D"  # noqa: S101, PLR2004
 

src/async_geotiff/_array.py

@@ -4,6 +4,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Self
 
+import async_tiff.enums
 from affine import Affine
 from async_tiff import TIFF
 from async_tiff.enums import PhotometricInterpretation
@@ -13,13 +14,14 @@
 from async_geotiff._ifd import IFDReference
 from async_geotiff._overview import Overview
 from async_geotiff._transform import TransformMixin
-from async_geotiff.enums import Compression, Interleaving
 
 if TYPE_CHECKING:
     from async_tiff import GeoKeyDirectory, ImageFileDirectory, ObspecInput
     from async_tiff.store import ObjectStore  # type: ignore # noqa: PGH003
     from pyproj import CRS
 
+    from async_geotiff.enums import Compression, Interleaving
+
 
 @dataclass(frozen=True, init=False, kw_only=True, repr=False)
 class GeoTIFF(FetchTileMixin, TransformMixin):
@@ -374,7 +376,7 @@ def has_geokeys(ifd: ImageFileDirectory) -> bool:
 def is_mask_ifd(ifd: ImageFileDirectory) -> bool:
     """Check if an IFD is a mask IFD."""
     return (
-        ifd.compression == Compression.deflate
+        ifd.compression == async_tiff.enums.CompressionMethod.Deflate
         and ifd.new_subfile_type is not None
         and ifd.new_subfile_type & 4 != 0
         and ifd.photometric_interpretation == PhotometricInterpretation.TransparencyMask

src/async_geotiff/_geotiff.py

@@ -34,6 +34,7 @@ def __call__(
         name: str,
         *,
         variant: Variant = "rasterio",
+        OVERVIEW_LEVEL: int | None = None,  # noqa: N803
         **kwargs: Any,  # noqa: ANN401
     ) -> Generator[DatasetReader, None, None]: ...
 
@@ -80,6 +81,7 @@ def _load(
         name: str,
         *,
         variant: Variant = "rasterio",
+        OVERVIEW_LEVEL: int | None = None,  # noqa: N803
         **kwargs: Any,  # noqa: ANN401
     ) -> Generator[DatasetReader, None, None]:
         path = f"{root_dir}/fixtures/geotiff-test-data/"
@@ -93,6 +95,10 @@ def _load(
             raise ValueError(f"Unknown variant: {variant}")
 
         path = f"{path}{name}.tif"
+
+        if OVERVIEW_LEVEL is not None and "OVERVIEW_LEVEL" not in kwargs:
+            kwargs["OVERVIEW_LEVEL"] = OVERVIEW_LEVEL
+
         with rasterio.open(path, **kwargs) as ds:
             yield ds
 

tests/conftest.py

@@ -79,3 +79,62 @@ async def test_fetch_overview(
 
     np.testing.assert_array_equal(tile.data, rasterio_data)
     assert tile.crs == geotiff.crs
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("file_name", "variant"),
+    [
+        ("maxar_opendata_yellowstone_visual", "vantor"),
+    ],
+)
+async def test_mask(
+    load_geotiff: LoadGeoTIFF,
+    load_rasterio: LoadRasterio,
+    file_name: str,
+    variant: Variant,
+) -> None:
+    geotiff = await load_geotiff(file_name, variant=variant)
+
+    tile = await geotiff.fetch_tile(0, 0)
+
+    assert tile.mask is not None
+    assert isinstance(tile.mask, np.ndarray)
+    assert tile.mask.dtype == np.bool_
+    assert tile.mask.shape == tile.data.shape[1:]
+
+    window = Window(0, 0, geotiff.tile_width, geotiff.tile_height)
+    with load_rasterio(file_name, variant=variant) as rasterio_ds:
+        mask = rasterio_ds.dataset_mask(window=window)
+
+    np.testing.assert_array_equal(tile.mask, mask.astype(np.bool_))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("file_name", "variant"),
+    [
+        ("maxar_opendata_yellowstone_visual", "vantor"),
+    ],
+)
+async def test_mask_overview(
+    load_geotiff: LoadGeoTIFF,
+    load_rasterio: LoadRasterio,
+    file_name: str,
+    variant: Variant,
+) -> None:
+    geotiff = await load_geotiff(file_name, variant=variant)
+    overview = geotiff.overviews[0]
+
+    tile = await overview.fetch_tile(0, 0)
+
+    assert tile.mask is not None
+    assert isinstance(tile.mask, np.ndarray)
+    assert tile.mask.dtype == np.bool_
+    assert tile.mask.shape == tile.data.shape[1:]
+
+    window = Window(0, 0, overview.tile_width, overview.tile_height)
+    with load_rasterio(file_name, variant=variant, OVERVIEW_LEVEL=0) as rasterio_ds:
+        mask = rasterio_ds.dataset_mask(window=window)
+
+    np.testing.assert_array_equal(tile.mask, mask.astype(np.bool_))