Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 41 additions & 7 deletions xrspatial/geotiff/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2683,9 +2683,19 @@ def read_geotiff_gpu(source: str, *,
# 5-8 today (the 2/3/4 fix in #1539 is in a sibling PR). Discard
# its geo_info and apply our own transform update below so the
# result is correct regardless of merge order.
#
# Forward ``max_pixels``, ``window``, and ``band`` so the
# caller's safety cap is honoured, windowed reads avoid
# decoding the full image, and single-band selection on a
# multi-band source skips the unused channels. Without this,
# the stripped GPU path bypassed all three (issue #1732).
# Orientation != 1 + window is already rejected at line 2495,
# so ``window`` is None whenever ``geo_info`` will be remapped
# below.
src.close()
arr_cpu, _ = _read_to_array(
source, overview_level=overview_level)
source, overview_level=overview_level,
window=window, band=band, max_pixels=max_pixels)
arr_gpu = cupy.asarray(arr_cpu)
if orientation != 1:
geo_info = _apply_orientation_geo_info(
Expand All @@ -2708,12 +2718,36 @@ 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)
# Apply window/band slicing post-decode. The stripped CPU
# fallback already produces the full-image array; slice on the
# GPU so the result matches ``open_geotiff`` /
# ``read_geotiff_dask`` semantics.
arr_gpu, coords = _gpu_apply_window_band(
arr_gpu, geo_info, window=window, band=band)
# ``read_to_array`` already applied window + band slicing, so
# ``arr_gpu`` is at output shape. Compute coords for that
# shape without re-slicing.
if window is not None:
r0, c0, r1, c1 = window
t = geo_info.transform
if t is None:
coords = {
'y': np.arange(r1 - r0, dtype=np.int64),
'x': np.arange(c1 - c0, dtype=np.int64),
}
elif geo_info.raster_type == RASTER_PIXEL_IS_POINT:
coords = {
'x': (np.arange(c0, c1, dtype=np.float64)
* t.pixel_width + t.origin_x),
'y': (np.arange(r0, r1, dtype=np.float64)
* t.pixel_height + t.origin_y),
}
else:
coords = {
'x': (np.arange(c0, c1, dtype=np.float64)
* t.pixel_width + t.origin_x
+ t.pixel_width * 0.5),
'y': (np.arange(r0, r1, dtype=np.float64)
* t.pixel_height + t.origin_y
+ t.pixel_height * 0.5),
}
else:
coords = _geo_to_coords(
geo_info, arr_gpu.shape[0], arr_gpu.shape[1])
# 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').
Expand Down
162 changes: 162 additions & 0 deletions xrspatial/geotiff/tests/test_gpu_stripped_forwarding_1732.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
"""Regression tests for issue #1732.

The stripped-TIFF fallback inside ``read_geotiff_gpu`` previously called
``read_to_array(source, overview_level=overview_level)`` and threw away
the caller's ``max_pixels``, ``window``, and ``band`` arguments. That
meant:

- a user-supplied ``max_pixels`` safety cap was silently ignored on
stripped files (the default ~1B pixel cap applied instead),
- windowed reads decoded the entire image before slicing on the GPU, and
- single-band selection on a multi-band stripped file still decoded
every band on the CPU.

These tests assert that all three kwargs are now forwarded to
``read_to_array`` so the stripped GPU path matches the contract of the
tiled GPU path.
"""
from __future__ import annotations

import importlib.util
import os
import tempfile

import numpy as np
import pytest
import xarray as xr


def _gpu_available() -> bool:
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",
)


@_gpu_only
def test_stripped_max_pixels_cap_is_enforced():
"""max_pixels smaller than the file must raise before full decode."""
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu

rng = np.random.RandomState(20260512)
data = rng.randint(0, 200, size=(64, 96)).astype(np.uint8)
da = xr.DataArray(data, dims=['y', 'x'])

with tempfile.TemporaryDirectory() as d:
p = os.path.join(d, 'tmp_1732_cap.tif')
to_geotiff(da, p, tiled=False)
# 64 * 96 = 6144 pixels; cap at 1000 must reject.
with pytest.raises(ValueError, match="max_pixels|pixel"):
read_geotiff_gpu(p, max_pixels=1000)


@_gpu_only
def test_stripped_window_returns_only_window():
"""Windowed read on a stripped file returns the window-sized array
with coords and transform that match the window origin.

The post-decode ``_gpu_apply_window_band`` call was replaced with a
coord-only computation in #1732. Compare against the CPU eager path
(which is the parity reference for this exact fixture) so a
regression in the coord-only branch -- or a drift in the windowed
``attrs['transform']`` -- shows up here.
"""
from xrspatial.geotiff import to_geotiff, open_geotiff, read_geotiff_gpu

rng = np.random.RandomState(20260512)
data = rng.randint(0, 200, size=(64, 96)).astype(np.uint8)
# Explicit y/x coords give the file a real georef so the coord-only
# path computes a non-trivial windowed transform / origin -- a plain
# ``dims=['y','x']`` array writes a no-georef TIFF where the coord
# branch degenerates to integer arange and would not catch a
# transform-math regression.
da = xr.DataArray(
data,
dims=['y', 'x'],
coords={
'y': np.arange(64, dtype=np.float64) * 0.5 + 100.0,
'x': np.arange(96, dtype=np.float64) * 0.5 + 200.0,
},
attrs={'crs': 4326},
)

with tempfile.TemporaryDirectory() as d:
p = os.path.join(d, 'tmp_1732_win.tif')
to_geotiff(da, p, tiled=False)
win = (8, 16, 40, 80) # 32x64 window
out = read_geotiff_gpu(p, window=win)
assert out.shape == (32, 64)
np.testing.assert_array_equal(out.data.get(), data[8:40, 16:80])
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in f087212: added explicit assertions on out.coords['y'], out.coords['x'], and out.attrs['transform'] so regressions in the coord-only path are caught.


# Coords + transform parity vs the CPU eager path. CPU runs
# through ``open_geotiff``'s own windowed-coord branch (line
# ~833), so any drift between the two coord computations is
# caught here.
cpu = open_geotiff(p, window=win)
np.testing.assert_array_equal(out.coords['y'].values,
cpu.coords['y'].values)
np.testing.assert_array_equal(out.coords['x'].values,
cpu.coords['x'].values)
# ``attrs['transform']`` carries the windowed origin (origin_x
# shifted by c0 * pixel_width, origin_y by r0 * pixel_height).
# Pin the exact tuple as well as parity with CPU so a regression
# in either ``_populate_attrs_from_geo_info`` or the coord-only
# branch is visible.
assert out.attrs['transform'] == cpu.attrs['transform']
# pixel_width=0.5, origin_x=200, c0=16 -> 200 + 16*0.5 = 208
# pixel_height=0.5, origin_y=100-0.5*0.5=99.75 (PixelIsArea),
# r0=8 -> 99.75 + 8*0.5 = 103.75
# to_geotiff writes the raw geo-transform (edge origin), so:
# origin_x_raw = 200 - 0.25 = 199.75; +16*0.5 = 207.75
# origin_y_raw = 100 - 0.25 = 99.75; +8*0.5 = 103.75
assert out.attrs['transform'] == (0.5, 0.0, 207.75,
0.0, 0.5, 103.75)


@_gpu_only
def test_stripped_band_selection_returns_2d():
"""Selecting band=1 on a 3-band stripped file returns a 2D array
matching the requested band."""
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu

rng = np.random.RandomState(20260512)
data = rng.randint(0, 200, size=(48, 80, 3)).astype(np.uint8)
da = xr.DataArray(data, dims=['y', 'x', 'band'])

with tempfile.TemporaryDirectory() as d:
p = os.path.join(d, 'tmp_1732_band.tif')
to_geotiff(da, p, tiled=False)
out = read_geotiff_gpu(p, band=1)
assert out.dims == ('y', 'x')
assert out.shape == (48, 80)
np.testing.assert_array_equal(out.data.get(), data[:, :, 1])


@_gpu_only
def test_stripped_window_plus_band():
"""Windowed read with band selection composes correctly."""
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu

rng = np.random.RandomState(20260512)
data = rng.randint(0, 200, size=(48, 80, 3)).astype(np.uint8)
da = xr.DataArray(data, dims=['y', 'x', 'band'])

with tempfile.TemporaryDirectory() as d:
p = os.path.join(d, 'tmp_1732_wb.tif')
to_geotiff(da, p, tiled=False)
win = (4, 8, 36, 72) # 32x64
out = read_geotiff_gpu(p, window=win, band=2)
assert out.dims == ('y', 'x')
assert out.shape == (32, 64)
np.testing.assert_array_equal(
out.data.get(), data[4:36, 8:72, 2])
Loading