Skip to content

Commit f9d1eb3

Browse files
authored
Add nvJPEG GPU acceleration for JPEG-compressed GeoTIFFs (#1050) (#1055)
Wire JPEG (TIFF tag 7) into the GPU decode and encode pipelines. When libnvjpeg.so is available, read_geotiff(gpu=True) and write_geotiff(gpu=True, compression='jpeg') use nvJPEG for hardware-accelerated JPEG on GPU. Falls back to Pillow on CPU when nvJPEG is not installed. Changes: - _gpu_decode.py: Add _find_nvjpeg_lib/_get_nvjpeg lazy discovery, _try_nvjpeg_batch_decode for GPU reads, _nvjpeg_batch_encode for GPU writes. Hook tag 7 into gpu_decode_tiles and gpu_compress_tiles. - _writer.py: Add 'jpeg' to _compression_tag map. Handle JPEG in _write_tiled and _write_stripped (call jpeg_compress directly, skip predictor). Validate uint8 dtype and 1/3 band constraint. - __init__.py: Update docstrings with JPEG option. - README.md: Add nvJPEG to GPU codecs list and JPEG write example. - tests/test_jpeg.py: 13 tests covering codec round trips, tag wiring, tiled/stripped write-read, public API, and dtype validation. - tests/test_writer.py: Fix test_unsupported_compression (was using 'jpeg' as the unsupported example).
1 parent 819ee71 commit f9d1eb3

File tree

6 files changed

+533
-11
lines changed

6 files changed

+533
-11
lines changed

README.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,12 +158,13 @@ read_geotiff('mosaic.vrt') # VRT mosaic (auto-detected
158158

159159
write_geotiff(cupy_array, 'out.tif') # auto-detects GPU
160160
write_geotiff(data, 'out.tif', gpu=True) # force GPU compress
161+
write_geotiff(data, 'ortho.tif', compression='jpeg') # JPEG for orthophotos
161162
write_vrt('mosaic.vrt', ['tile1.tif', 'tile2.tif']) # generate VRT
162163
```
163164

164165
**Compression codecs:** Deflate, LZW (Numba JIT), ZSTD, PackBits, JPEG (Pillow), uncompressed
165166

166-
**GPU codecs:** Deflate and ZSTD via nvCOMP batch API; LZW via Numba CUDA kernels
167+
**GPU codecs:** Deflate and ZSTD via nvCOMP; LZW via Numba CUDA; JPEG via nvJPEG
167168

168169
**Features:**
169170
- Tiled, stripped, BigTIFF, multi-band (RGB/RGBA), sub-byte (1/2/4/12-bit)

xrspatial/geotiff/__init__.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,10 @@ def write_geotiff(data: xr.DataArray | np.ndarray, path: str, *,
319319
nodata : float, int, or None
320320
NoData value.
321321
compression : str
322-
'none', 'deflate', or 'lzw'.
322+
'none', 'deflate', 'lzw', 'jpeg', 'packbits', or 'zstd'.
323+
JPEG is lossy and only supports uint8 data (1 or 3 bands).
324+
With ``gpu=True``, JPEG uses nvJPEG for GPU-accelerated
325+
encode/decode when available, falling back to Pillow on CPU.
323326
tiled : bool
324327
Use tiled layout (default True).
325328
tile_size : int
@@ -756,7 +759,8 @@ def write_geotiff_gpu(data, path: str, *,
756759
nodata : float, int, or None
757760
NoData value.
758761
compression : str
759-
'zstd' (default, fastest on GPU), 'deflate', or 'none'.
762+
'zstd' (default, fastest on GPU), 'deflate', 'jpeg', or 'none'.
763+
JPEG uses nvJPEG when available, falling back to Pillow.
760764
tile_size : int
761765
Tile size in pixels (default 256).
762766
predictor : bool

0 commit comments

Comments
 (0)