Skip to content

Commit 7637a42

Browse files
committed
geotiff: forward max_pixels, window, band on GPU stripped fallback (#1732)
The stripped-TIFF branch of read_geotiff_gpu called read_to_array(source, overview_level=overview_level), dropping the caller's max_pixels, window, and band kwargs. Three consequences: - max_pixels safety cap bypassed (default 1B pixel limit applied instead of the smaller value the caller set). - Windowed reads decoded the entire image before slicing on the GPU. - Single-band reads on multi-band sources decoded every band. Forward all three. The post-decode _gpu_apply_window_band call is replaced with coord-only computation since read_to_array now produces the pre-sliced array. Adds regression coverage in test_gpu_stripped_forwarding_1732.py.
1 parent 1624d13 commit 7637a42

2 files changed

Lines changed: 159 additions & 7 deletions

File tree

xrspatial/geotiff/__init__.py

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2529,9 +2529,19 @@ def read_geotiff_gpu(source: str, *,
25292529
# 5-8 today (the 2/3/4 fix in #1539 is in a sibling PR). Discard
25302530
# its geo_info and apply our own transform update below so the
25312531
# result is correct regardless of merge order.
2532+
#
2533+
# Forward ``max_pixels``, ``window``, and ``band`` so the
2534+
# caller's safety cap is honoured, windowed reads avoid
2535+
# decoding the full image, and single-band selection on a
2536+
# multi-band source skips the unused channels. Without this,
2537+
# the stripped GPU path bypassed all three (issue #1732).
2538+
# Orientation != 1 + window is already rejected at line 2495,
2539+
# so ``window`` is None whenever ``geo_info`` will be remapped
2540+
# below.
25322541
src.close()
25332542
arr_cpu, _ = read_to_array(
2534-
source, overview_level=overview_level)
2543+
source, overview_level=overview_level,
2544+
window=window, band=band, max_pixels=max_pixels)
25352545
arr_gpu = cupy.asarray(arr_cpu)
25362546
if orientation != 1:
25372547
geo_info = _apply_orientation_geo_info(
@@ -2554,12 +2564,36 @@ def read_geotiff_gpu(source: str, *,
25542564
target = np.dtype(dtype)
25552565
_validate_dtype_cast(np.dtype(str(arr_gpu.dtype)), target)
25562566
arr_gpu = arr_gpu.astype(target)
2557-
# Apply window/band slicing post-decode. The stripped CPU
2558-
# fallback already produces the full-image array; slice on the
2559-
# GPU so the result matches ``open_geotiff`` /
2560-
# ``read_geotiff_dask`` semantics.
2561-
arr_gpu, coords = _gpu_apply_window_band(
2562-
arr_gpu, geo_info, window=window, band=band)
2567+
# ``read_to_array`` already applied window + band slicing, so
2568+
# ``arr_gpu`` is at output shape. Compute coords for that
2569+
# shape without re-slicing.
2570+
if window is not None:
2571+
r0, c0, r1, c1 = window
2572+
t = geo_info.transform
2573+
if t is None:
2574+
coords = {
2575+
'y': np.arange(r1 - r0, dtype=np.int64),
2576+
'x': np.arange(c1 - c0, dtype=np.int64),
2577+
}
2578+
elif geo_info.raster_type == RASTER_PIXEL_IS_POINT:
2579+
coords = {
2580+
'x': (np.arange(c0, c1, dtype=np.float64)
2581+
* t.pixel_width + t.origin_x),
2582+
'y': (np.arange(r0, r1, dtype=np.float64)
2583+
* t.pixel_height + t.origin_y),
2584+
}
2585+
else:
2586+
coords = {
2587+
'x': (np.arange(c0, c1, dtype=np.float64)
2588+
* t.pixel_width + t.origin_x
2589+
+ t.pixel_width * 0.5),
2590+
'y': (np.arange(r0, r1, dtype=np.float64)
2591+
* t.pixel_height + t.origin_y
2592+
+ t.pixel_height * 0.5),
2593+
}
2594+
else:
2595+
coords = _geo_to_coords(
2596+
geo_info, arr_gpu.shape[0], arr_gpu.shape[1])
25632597
# Multi-band stripped reads come back as (y, x, band); mirror
25642598
# the tiled branch so dims line up with ndim. Single-band stays
25652599
# 2-D ('y', 'x').
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
"""Regression tests for issue #1732.
2+
3+
The stripped-TIFF fallback inside ``read_geotiff_gpu`` previously called
4+
``read_to_array(source, overview_level=overview_level)`` and threw away
5+
the caller's ``max_pixels``, ``window``, and ``band`` arguments. That
6+
meant:
7+
8+
- a user-supplied ``max_pixels`` safety cap was silently ignored on
9+
stripped files (the default ~1B pixel cap applied instead),
10+
- windowed reads decoded the entire image before slicing on the GPU, and
11+
- single-band selection on a multi-band stripped file still decoded
12+
every band on the CPU.
13+
14+
These tests assert that all three kwargs are now forwarded to
15+
``read_to_array`` so the stripped GPU path matches the contract of the
16+
tiled GPU path.
17+
"""
18+
from __future__ import annotations
19+
20+
import importlib.util
21+
import os
22+
import tempfile
23+
24+
import numpy as np
25+
import pytest
26+
import xarray as xr
27+
28+
29+
def _gpu_available() -> bool:
30+
if importlib.util.find_spec("cupy") is None:
31+
return False
32+
try:
33+
import cupy
34+
return bool(cupy.cuda.is_available())
35+
except Exception:
36+
return False
37+
38+
39+
_HAS_GPU = _gpu_available()
40+
_gpu_only = pytest.mark.skipif(
41+
not _HAS_GPU,
42+
reason="cupy + CUDA required",
43+
)
44+
45+
46+
@_gpu_only
47+
def test_stripped_max_pixels_cap_is_enforced():
48+
"""max_pixels smaller than the file must raise before full decode."""
49+
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu
50+
51+
rng = np.random.RandomState(20260512)
52+
data = rng.randint(0, 200, size=(64, 96)).astype(np.uint8)
53+
da = xr.DataArray(data, dims=['y', 'x'])
54+
55+
with tempfile.TemporaryDirectory() as d:
56+
p = os.path.join(d, 'tmp_1732_cap.tif')
57+
to_geotiff(da, p, tiled=False)
58+
# 64 * 96 = 6144 pixels; cap at 1000 must reject.
59+
with pytest.raises(ValueError, match="max_pixels|pixel"):
60+
read_geotiff_gpu(p, max_pixels=1000)
61+
62+
63+
@_gpu_only
64+
def test_stripped_window_returns_only_window():
65+
"""Windowed read on a stripped file returns the window-sized array
66+
with coords that match the window origin."""
67+
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu
68+
69+
rng = np.random.RandomState(20260512)
70+
data = rng.randint(0, 200, size=(64, 96)).astype(np.uint8)
71+
da = xr.DataArray(data, dims=['y', 'x'])
72+
73+
with tempfile.TemporaryDirectory() as d:
74+
p = os.path.join(d, 'tmp_1732_win.tif')
75+
to_geotiff(da, p, tiled=False)
76+
win = (8, 16, 40, 80) # 32x64 window
77+
out = read_geotiff_gpu(p, window=win)
78+
assert out.shape == (32, 64)
79+
np.testing.assert_array_equal(out.data.get(), data[8:40, 16:80])
80+
81+
82+
@_gpu_only
83+
def test_stripped_band_selection_returns_2d():
84+
"""Selecting band=1 on a 3-band stripped file returns a 2D array
85+
matching the requested band."""
86+
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu
87+
88+
rng = np.random.RandomState(20260512)
89+
data = rng.randint(0, 200, size=(48, 80, 3)).astype(np.uint8)
90+
da = xr.DataArray(data, dims=['y', 'x', 'band'])
91+
92+
with tempfile.TemporaryDirectory() as d:
93+
p = os.path.join(d, 'tmp_1732_band.tif')
94+
to_geotiff(da, p, tiled=False)
95+
out = read_geotiff_gpu(p, band=1)
96+
assert out.dims == ('y', 'x')
97+
assert out.shape == (48, 80)
98+
np.testing.assert_array_equal(out.data.get(), data[:, :, 1])
99+
100+
101+
@_gpu_only
102+
def test_stripped_window_plus_band():
103+
"""Windowed read with band selection composes correctly."""
104+
from xrspatial.geotiff import to_geotiff, read_geotiff_gpu
105+
106+
rng = np.random.RandomState(20260512)
107+
data = rng.randint(0, 200, size=(48, 80, 3)).astype(np.uint8)
108+
da = xr.DataArray(data, dims=['y', 'x', 'band'])
109+
110+
with tempfile.TemporaryDirectory() as d:
111+
p = os.path.join(d, 'tmp_1732_wb.tif')
112+
to_geotiff(da, p, tiled=False)
113+
win = (4, 8, 36, 72) # 32x64
114+
out = read_geotiff_gpu(p, window=win, band=2)
115+
assert out.dims == ('y', 'x')
116+
assert out.shape == (32, 64)
117+
np.testing.assert_array_equal(
118+
out.data.get(), data[4:36, 8:72, 2])

0 commit comments

Comments
 (0)