Handle PlanarConfiguration=2 across CPU + GPU multi-band reads

xrspatial/geotiff/__init__.py

@@ -13,6 +13,7 @@
 """
 from __future__ import annotations
 
+import math
 import warnings
 
 import numpy as np
@@ -1396,6 +1397,92 @@ def _read():
     return _read()
 
 
+def _gpu_decode_single_band_tiles(
+    source, offsets, byte_counts,
+    tw, th, width, height,
+    compression, predictor, file_dtype,
+    *,
+    byte_order: str,
+    gpu: str,
+    overview_level,
+    cupy,
+):
+    """Decode one band's tile sequence into a 2-D ``(H, W)`` cupy array.
+
+    Helper for the ``PlanarConfiguration=2`` GPU path: the existing
+    ``gpu_decode_tiles`` / ``gpu_decode_tiles_from_file`` kernels assume
+    a single chunky tile sequence with ``bytes_per_pixel = itemsize *
+    samples``. For planar=2 each band has its own list of tiles, so we
+    invoke those kernels once per band with ``samples=1`` and stack the
+    resulting 2-D arrays into ``(H, W, samples)`` afterwards.
+
+    Mirrors the two-stage GPU pipeline in ``read_geotiff_gpu`` -- GDS
+    first, then CPU mmap + GPU decode -- and honours the same
+    ``gpu='strict'`` / ``'auto'`` semantics. Sparse tiles are not
+    expected here; the caller routes sparse files to the CPU path.
+    """
+    from ._reader import _FileSource
+    from ._gpu_decode import gpu_decode_tiles, gpu_decode_tiles_from_file
+
+    arr_gpu = None
+    try:
+        arr_gpu = gpu_decode_tiles_from_file(
+            source, offsets, byte_counts,
+            tw, th, width, height,
+            compression, predictor, file_dtype, 1,
+            byte_order=byte_order,
+        )
+    except Exception as e:
+        if gpu == 'strict':
+            raise
+        warnings.warn(
+            f"read_geotiff_gpu: GPU decode failed "
+            f"({type(e).__name__}: {e}); falling back to CPU.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+        arr_gpu = None
+
+    if arr_gpu is None:
+        src2 = _FileSource(source)
+        data2 = src2.read_all()
+        try:
+            compressed_tiles = [
+                bytes(data2[offsets[i]:offsets[i] + byte_counts[i]])
+                for i in range(len(offsets))
+            ]
+        finally:
+            src2.close()
+        try:
+            arr_gpu = gpu_decode_tiles(
+                compressed_tiles,
+                tw, th, width, height,
+                compression, predictor, file_dtype, 1,
+                byte_order=byte_order,
+            )
+        except Exception as e:
+            if gpu == 'strict':
+                raise
+            warnings.warn(
+                f"read_geotiff_gpu: GPU decode failed "
+                f"({type(e).__name__}: {e}); falling back to CPU.",
+                RuntimeWarning,
+                stacklevel=3,
+            )
+            # Per-band CPU fallback would require a single-band CPU
+            # decode path keyed off planar config; the easier safe move
+            # is to surface the failure so the caller upstream can fall
+            # back to read_to_array for the whole image.
+            raise
+
+    if arr_gpu.ndim != 2 or arr_gpu.shape != (height, width):
+        raise RuntimeError(
+            f"single-band GPU tile decode produced shape "
+            f"{arr_gpu.shape}, expected ({height}, {width})"
+        )
+    return arr_gpu
+
+
 def read_geotiff_gpu(source: str, *,
                      dtype=None,
                      overview_level: int | None = None,
@@ -1519,14 +1606,24 @@ def read_geotiff_gpu(source: str, *,
                 target = np.dtype(dtype)
                 _validate_dtype_cast(np.dtype(str(arr_gpu.dtype)), target)
                 arr_gpu = arr_gpu.astype(target)
-            return xr.DataArray(arr_gpu, dims=['y', 'x'],
+            # Multi-band stripped reads come back as (y, x, band); mirror
+            # the tiled branch so dims line up with ndim. Single-band stays
+            # 2-D ('y', 'x'). Same handling as the GPU stripped fix in
+            # PR #1525.
+            if arr_gpu.ndim == 3:
+                dims = ['y', 'x', 'band']
+                coords['band'] = np.arange(arr_gpu.shape[2])
+            else:
+                dims = ['y', 'x']
+            return xr.DataArray(arr_gpu, dims=dims,
                                 coords=coords, name=name, attrs=attrs)
 
         offsets = ifd.tile_offsets
         byte_counts = ifd.tile_byte_counts
         compression = ifd.compression
         predictor = ifd.predictor
         samples = ifd.samples_per_pixel
+        planar = ifd.planar_config
         tw = ifd.tile_width
         th = ifd.tile_height
         width = ifd.width
@@ -1552,7 +1649,47 @@ def read_geotiff_gpu(source: str, *,
     # Sparse tiles (byte_count == 0) are unsupported on the GPU pipeline;
     # the CPU reader fills them with nodata and copies onto the GPU.
     has_sparse_tile = any(bc == 0 for bc in byte_counts)
-    if has_sparse_tile:
+
+    # PlanarConfiguration=2 (separate bands): each band has its own list
+    # of tiles back-to-back in TileOffsets / TileByteCounts. The GPU
+    # tile-assembly kernel assumes a single chunky tile sequence with
+    # bytes_per_pixel = itemsize * samples, so it cannot handle planar=2
+    # directly. Decode each band's tile slab as a single-band image, then
+    # stack into (H, W, samples). For planar=1 (chunky) the existing
+    # single-pass kernel is correct.
+    if planar == 2 and samples > 1:
+        tiles_across = math.ceil(width / tw)
+        tiles_down = math.ceil(height / th)
+        tiles_per_band = tiles_across * tiles_down
+        if tiles_per_band * samples != len(offsets):
+            raise ValueError(
+                f"PlanarConfiguration=2 expects "
+                f"{tiles_per_band * samples} TileOffsets entries "
+                f"({tiles_across} x {tiles_down} x {samples} bands), "
+                f"got {len(offsets)}"
+            )
+        band_arrays = []
+        for band_idx in range(samples):
+            b0 = band_idx * tiles_per_band
+            b1 = b0 + tiles_per_band
+            band_offsets = list(offsets[b0:b1])
+            band_byte_counts = list(byte_counts[b0:b1])
+            band_arr = _gpu_decode_single_band_tiles(
+                source, band_offsets, band_byte_counts,
+                tw, th, width, height,
+                compression, predictor, file_dtype,
+                byte_order=header.byte_order,
+                gpu=gpu, overview_level=overview_level,
+                cupy=cupy,
+            )
+            band_arrays.append(band_arr)
+        arr_gpu = cupy.stack(band_arrays, axis=2)
+        # Sanity check: assembled shape must match (H, W, samples).
+        assert arr_gpu.shape == (height, width, samples), (
+            f"planar=2 GPU assembly produced shape {arr_gpu.shape}, "
+            f"expected ({height}, {width}, {samples})"
+        )
+    elif has_sparse_tile:
         arr_cpu, _ = read_to_array(source, overview_level=overview_level)
         arr_gpu = cupy.asarray(arr_cpu)
     else:
@@ -1609,6 +1746,15 @@ def read_geotiff_gpu(source: str, *,
             arr_cpu, _ = read_to_array(source, overview_level=overview_level)
             arr_gpu = cupy.asarray(arr_cpu)
 
+    # Multi-band tiled output must be (H, W, samples) regardless of planar
+    # config -- catch any shape regression in the kernels before we attach
+    # dims/coords below.
+    if samples > 1:
+        assert arr_gpu.shape[:2] == (height, width) and arr_gpu.shape[2] == samples, (
+            f"GPU multi-band tile assembly produced shape {arr_gpu.shape}, "
+            f"expected ({height}, {width}, {samples})"
+        )
+
     if dtype is not None:
         target = np.dtype(dtype)
         _validate_dtype_cast(np.dtype(str(arr_gpu.dtype)), target)

xrspatial/geotiff/tests/test_planar_multiband.py

@@ -0,0 +1,195 @@
+"""Tests for ``PlanarConfiguration`` handling in CPU and GPU multi-band reads.
+
+Covers two bugs surfaced by the geotiff accuracy audit:
+
+A2 (HIGH)
+    GPU tiled multi-band reads returned wrong pixel values because the
+    tile-assembly kernel iterated a single ``tile_offsets`` list with
+    ``bytes_per_pixel = itemsize * samples``. Planar=2 files (one tile
+    sequence per band, contiguous in TileOffsets) silently produced
+    garbage because the kernel had no notion of separate band slabs.
+
+A3 (MEDIUM)
+    The CPU stripped reader returned planar=2 multi-band shape with
+    axes mismatched against the assigned dims. Already fixed in the CPU
+    path; this module pins the contract so future refactors do not
+    regress it.
+
+Each combination -- planar config x layout x band count x dtype --
+round-trips through the writer (or tifffile, where we want to control
+planar config and layout independently) and asserts that the CPU read,
+the GPU read, and the original numpy buffer all agree pixel-for-pixel.
+
+GPU tests skip cleanly when CUDA is unavailable.
+"""
+from __future__ import annotations
+
+import importlib.util
+import os
+import tempfile
+
+import numpy as np
+import pytest
+
+
+tifffile = pytest.importorskip("tifffile")
+
+
+def _gpu_available() -> bool:
+    """True if cupy is importable and CUDA is initialised."""
+    if importlib.util.find_spec("cupy") is None:
+        return False
+    try:
+        import cupy
+
+        return bool(cupy.cuda.is_available())
+    except Exception:
+        return False
+
+
+_HAS_GPU = _gpu_available()
+_gpu_only = pytest.mark.skipif(
+    not _HAS_GPU,
+    reason="cupy + CUDA required",
+)
+
+
+def _make_data(bands: int, height: int, width: int, dtype) -> np.ndarray:
+    """Deterministic random multi-band raster shaped (bands, height, width)."""
+    rng = np.random.RandomState(0xA2A3 + bands * 1000 + height + np.dtype(dtype).itemsize)
+    info = np.iinfo(dtype)
+    high = min(int(info.max), 60_000) + 1
+    return rng.randint(0, high, size=(bands, height, width)).astype(dtype)
+
+
+def _write_tiff(path: str, data: np.ndarray, *, planar: str, tiled: bool):
+    """Write *data* (bands, height, width) with the requested planar config + layout.
+
+    tifffile defaults: when bands are first axis, planar='separate'.
+    Pass planar='contig' to interleave samples (data is transposed to
+    (height, width, bands) before write).
+    """
+    kwargs = {"photometric": "rgb" if data.shape[0] == 3 else "minisblack"}
+    if data.shape[0] not in (1, 3):
+        # tifffile rejects 'rgb' for non-3-band; use a generic photometric
+        # tag and tell it to skip extra-sample warnings.
+        kwargs = {"photometric": "minisblack"}
+    if tiled:
+        kwargs["tile"] = (32, 32)
+    if planar == "separate":
+        kwargs["planarconfig"] = "separate"
+        tifffile.imwrite(path, data, **kwargs)
+    elif planar == "contig":
+        kwargs["planarconfig"] = "contig"
+        tifffile.imwrite(path, np.transpose(data, (1, 2, 0)), **kwargs)
+    else:
+        raise ValueError(f"unknown planar={planar!r}")
+
+
+@pytest.mark.parametrize("planar", ["separate", "contig"])
+@pytest.mark.parametrize("tiled", [True, False])
+@pytest.mark.parametrize("bands", [2, 3, 4])
+@pytest.mark.parametrize("dtype", [np.uint8, np.uint16])
+def test_cpu_planar_multiband(planar, tiled, bands, dtype):
+    """CPU ``open_geotiff`` returns (y, x, band) for every planar+layout combo."""
+    from xrspatial.geotiff import open_geotiff
+
+    height, width = 64, 96
+    data = _make_data(bands, height, width, dtype)
+    expected = np.transpose(data, (1, 2, 0))
+
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "x.tif")
+        _write_tiff(path, data, planar=planar, tiled=tiled)
+        out = open_geotiff(path)
+
+    assert out.dims == ("y", "x", "band"), (
+        f"got dims {out.dims} for planar={planar} tiled={tiled} "
+        f"bands={bands} dtype={np.dtype(dtype).name}"
+    )
+    assert out.shape == (height, width, bands)
+    arr = np.asarray(out.data)
+    assert arr.dtype == np.dtype(dtype)
+    np.testing.assert_array_equal(arr, expected)
+
+
+@_gpu_only
+@pytest.mark.parametrize("planar", ["separate", "contig"])
+@pytest.mark.parametrize("tiled", [True, False])
+@pytest.mark.parametrize("bands", [2, 3, 4])
+@pytest.mark.parametrize("dtype", [np.uint8, np.uint16])
+def test_gpu_matches_cpu_planar_multiband(planar, tiled, bands, dtype):
+    """GPU read agrees with CPU read AND with the source array."""
+    from xrspatial.geotiff import open_geotiff, read_geotiff_gpu
+
+    height, width = 64, 96
+    data = _make_data(bands, height, width, dtype)
+    expected = np.transpose(data, (1, 2, 0))
+
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "x.tif")
+        _write_tiff(path, data, planar=planar, tiled=tiled)
+        cpu = np.asarray(open_geotiff(path).data)
+        gpu_da = read_geotiff_gpu(path)
+
+    assert gpu_da.dims == ("y", "x", "band")
+    assert gpu_da.shape == (height, width, bands)
+    gpu = gpu_da.data.get()
+
+    np.testing.assert_array_equal(cpu, expected)
+    np.testing.assert_array_equal(gpu, expected)
+    np.testing.assert_array_equal(gpu, cpu)
+
+
+@pytest.mark.parametrize("tiled", [True, False])
+def test_cpu_singleband_sanity(tiled):
+    """Single-band reads stay 2-D regardless of layout."""
+    from xrspatial.geotiff import open_geotiff
+
+    rng = np.random.RandomState(7)
+    data = rng.randint(0, 200, size=(48, 80)).astype(np.uint8)
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "s.tif")
+        kwargs = {"photometric": "minisblack"}
+        if tiled:
+            kwargs["tile"] = (32, 32)
+        tifffile.imwrite(path, data, **kwargs)
+        out = open_geotiff(path)
+    assert out.dims == ("y", "x")
+    assert out.shape == (48, 80)
+    np.testing.assert_array_equal(np.asarray(out.data), data)
+
+
+@_gpu_only
+@pytest.mark.parametrize("tiled", [True, False])
+def test_gpu_singleband_sanity(tiled):
+    """Single-band GPU reads stay 2-D regardless of layout."""
+    from xrspatial.geotiff import read_geotiff_gpu
+
+    rng = np.random.RandomState(7)
+    data = rng.randint(0, 200, size=(48, 80)).astype(np.uint8)
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "s.tif")
+        kwargs = {"photometric": "minisblack"}
+        if tiled:
+            kwargs["tile"] = (32, 32)
+        tifffile.imwrite(path, data, **kwargs)
+        out = read_geotiff_gpu(path)
+    assert out.dims == ("y", "x")
+    assert out.shape == (48, 80)
+    np.testing.assert_array_equal(out.data.get(), data)
+
+
+def test_a3_repro_stripped_planar_separate():
+    """Spec-level guard for A3: stripped planar=2 must yield (y, x, band)."""
+    from xrspatial.geotiff import open_geotiff
+
+    data = _make_data(3, 64, 96, np.uint8)
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "sp.tif")
+        tifffile.imwrite(path, data, photometric="rgb")  # default planar=2
+        out = open_geotiff(path)
+    assert out.dims == ("y", "x", "band")
+    assert out.shape == (64, 96, 3), (
+        f"got shape {out.shape} for dims {out.dims} -- A3 regressed"
+    )