API Reference⚓︎

Get image pixel data from GeoTIFF as a DLPack capsule.

Parameters:

• stream ⚓︎

  (int | None, default: None ) –

  A Python integer representing a pointer to a stream, on devices that support streams. Device-specific values of stream for CUDA:

  • None: producer must assume the legacy default stream (default).
  • 1: the legacy default stream.
  • 2: the per-thread default stream.
  • > 2: stream number represented as a Python integer.
  • 0 is disallowed due to its ambiguity: 0 could mean either None, 1, or 2.

• max_version ⚓︎

  (tuple[int, int] | None, default: None ) –

  The maximum DLPack version that the consumer (i.e., the caller of __dlpack__) supports, in the form of a 2-tuple (major, minor). This method may return a capsule of version max_version (recommended if it does support that), or of a different version. This means the consumer must verify the version even when max_version is passed.

• dl_device ⚓︎

  (tuple[int, int] | None, default: None ) –

  The DLPack device type. Default is None, meaning the exported capsule should be on the same device as self is (i.e. CUDA). When specified, the format must be a 2-tuple, following that of the return value of array.__dlpack_device__(). If the device type cannot be handled by the producer, this function will raise BufferError.

• copy ⚓︎

  (bool | None, default: None ) –

  Boolean indicating whether or not to copy the input. Currently only None is supported, meaning the function must reuse the existing memory buffer if possible and copy otherwise (copy is not actually implemented).

Returns:

• tensor ( CapsuleType ) –

  1D tensor in row-major order containing the GeoTIFF pixel data.

Raises:

• NotImplementedError –

  If either of these cases happen:

  • stream>2 is passed in, as only legacy default stream (1) or per-thread default stream (2) is supported for now.
  • max_version is incompatible with the DLPack major version in this library.
  • copy is set to a value other than None as Copy keyword argument behavior is not handled yet.
• BufferError –

  If trying to decode to non-CUDA memory, i.e. when dl_device is not None, or set to a tuple other than (2, x). This error may also be raised if trying to decode to an unsupported version from the DLPack 0.x series.

Get device type and device ID in DLPack format.

Meant for use by from_dlpack().

Returns:

• device ( (int, int) ) –

  A tuple (device_type, device_id) in DLPack format.

Get list of x and y coordinates.

Determined based on an Affine transformation matrix built from the ModelPixelScaleTag and ModelTiepointTag TIFF tags. Note that non-zero rotation (set by ModelTransformationTag) is currently unsupported.

Returns:

• coords ( (ndarray, ndarray) ) –

  A tuple (x_coords, y_coords) of numpy.ndarray objects representing the GeoTIFF's x- and y-coordinates.

Get image pixel data from GeoTIFF as a DLPack capsule

Returns:

• tensor ( CapsuleType ) –

  3D tensor of shape (band, height, width) containing the GeoTIFF pixel data.

Raises:

• NotImplementedError –

  If stream is not None, as only decoding to the CPU is supported.

Get device type and device ID in DLPack format.

Meant for use by from_dlpack().

Returns:

• device ( (int, int) ) –

  A tuple (device_type, device_id) in DLPack format.

Get list of x and y coordinates.

Determined based on an Affine transformation matrix built from the ModelPixelScaleTag and ModelTiepointTag TIFF tags. Note that non-zero rotation (set by ModelTransformationTag) is currently unsupported.

Returns:

• coords ( (ndarray, ndarray) ) –

  A tuple (x_coords, y_coords) of numpy.ndarray objects representing the GeoTIFF's x- and y-coordinates.

Backend open_dataset method used by Xarray in xarray.open_dataset.

Parameters:

• filename_or_obj ⚓︎

  (str) –

  File path or url to a TIFF (.tif) image file that can be read by the nvTIFF or image-tiff backend library.

• device ⚓︎

  ((int, int) | None, default: (2, 0) ) –

  Device on which to place the created array, given as a tuple in the form of (device_type, device_id), corresponding to DLPack's device convention. Default is (2, 0) which means device_type='cuda', device_id=0. Pass device=None or device=(1, 0) to use the CPU fallback.

Returns:

• Dataset –

 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

def open_dataset(  # type: ignore[override]
    self,
    filename_or_obj: str,
    *,
    drop_variables: str | Iterable[str] | None = None,
    device: tuple[int, int] | None = (2, 0),
    # other backend specific keyword arguments
    # `chunks` and `cache` DO NOT go here, they are handled by xarray
    mask_and_scale=None,
) -> xr.Dataset:
    """
    Backend open_dataset method used by Xarray in [xarray.open_dataset][].

    Parameters
    ----------
    filename_or_obj : str
        File path or url to a TIFF (.tif) image file that can be read by the
        nvTIFF or image-tiff backend library.
    device : (int, int) | None
        Device on which to place the created array, given as a tuple in the form of
        (device_type, device_id), corresponding to
        [DLPack's device convention](https://data-apis.org/array-api/2024.12/API_specification/generated/array_api.array.__dlpack_device__.html).
        Default is (2, 0) which means device_type='cuda', device_id=0. Pass
        `device=None` or `device=(1, 0)` to use the CPU fallback.

    Returns
    -------
    xarray.Dataset

    """
    if device is None:
        device = (1, 0)  # CPU device fallback
    # TODO handle device.__dlpack_device__ ?
    # https://github.com/data-apis/array-api/issues/972
    device_type, device_id = device
    match device_type:
        case 1:  # CPU (image-tiff backend)
            cog = CogReader(path=filename_or_obj)
            array: np.ndarray = np.from_dlpack(cog)
            x_coords, y_coords = cog.xy_coords()  # TODO consider using rasterix
            channels, _height, _width = array.shape
        case 2:  # CUDA (nvTIFF backend)
            import cupy as cp  # type: ignore[import-untyped]

            from cog3pio import CudaCogReader

            with cp.cuda.Stream(ptds=True):
                cog_ = CudaCogReader(path=filename_or_obj, device_id=device_id)
                array_: cp.ndarray = cp.from_dlpack(cog_)  # 1-D Array
                x_coords, y_coords = cog_.xy_coords()  # TODO consider rasterix
                height, width = (len(y_coords), len(x_coords))
                channels: int = len(array_) // (height * width)  # type: ignore[no-redef]
                # TODO make API to get proper 3-D shape directly, or use cuTENSOR
                array_ = array_.reshape(height, width, channels)  # HWC
                array = array_.transpose(2, 0, 1)  # CHW
        case _:
            msg = (
                "Currently only support decoding to DLPack device_type 2 (CUDA) or "
                f"1 (CPU), but got {device_type}"
            )
            raise NotImplementedError(msg)

    dataarray: xr.DataArray = xr.DataArray(
        data=array,
        coords={
            "band": np.arange(channels, dtype=np.uint8),
            "y": y_coords,
            "x": x_coords,
        },
        name=None,
        attrs=None,
    )

    return dataarray.to_dataset(name="raster")