buteo.raster.shift

  1"""### Shift rasters. ###
  2
  3Module to shift the location of rasters in geographic coordinates.
  4"""
  5
  6# Standard library
  7from typing import Union, Optional, List
  8
  9# External
 10import numpy as np
 11from osgeo import gdal
 12
 13# Internal
 14from buteo.utils import (
 15    utils_base,
 16    utils_gdal,
 17    utils_path,
 18    utils_io,
 19)
 20from buteo.raster import core_raster, core_raster_io
 21from buteo.array.convolution import convolve_array_simple
 22from buteo.array.convolution_kernels import kernel_shift
 23
 24
 25
 26def _raster_shift(
 27    raster: Union[str, gdal.Dataset],
 28    shift_list: List[Union[int, float]],
 29    out_path: Optional[str] = None,
 30    *,
 31    overwrite: bool = True,
 32    creation_options: Optional[List[str]] = None,
 33) -> Union[str, gdal.Dataset]:
 34    """Internal."""
 35    assert isinstance(shift_list, (list, tuple)), f"shift_list must be a list or a tuple. {shift_list}"
 36    assert len(shift_list) == 2, f"shift_list must be a list or tuple with len 2 (x_shift, y_shift): {shift_list}"
 37
 38    for shift in shift_list:
 39        assert isinstance(shift, (int, float)), f"shift must be an int or a float: {shift}"
 40
 41    ref = core_raster._raster_open(raster)
 42    metadata = core_raster._get_basic_metadata_raster(ref)
 43
 44    x_shift, y_shift = shift_list
 45
 46    if out_path is None:
 47        out_path = utils_path._get_temp_filepath("shifted_raster.tif")
 48    else:
 49        if not utils_path._check_is_valid_output_filepath(out_path, overwrite=overwrite):
 50            raise ValueError(f"out_path is not a valid output path: {out_path}")
 51
 52    utils_path._delete_if_required(out_path, overwrite)
 53
 54    driver = gdal.GetDriverByName(utils_gdal._get_raster_driver_name_from_path(out_path))
 55
 56    shifted = driver.Create(
 57        out_path,  # Location of the saved raster, ignored if driver is memory.
 58        metadata["width"],  # Dataframe width in pixels (e.g. 1920px).
 59        metadata["height"],  # Dataframe height in pixels (e.g. 1280px).
 60        metadata["bands"],  # The number of bands required.
 61        metadata["dtype_gdal"],  # Datatype of the destination.
 62        utils_gdal._get_default_creation_options(creation_options),
 63    )
 64
 65    new_transform = list(metadata["geotransform"])
 66    new_transform[0] += x_shift
 67    new_transform[3] += y_shift
 68
 69    shifted.SetGeoTransform(new_transform)
 70    shifted.SetProjection(metadata["projection_wkt"])
 71
 72    src_nodata = metadata["nodata_value"]
 73
 74    for band in range(metadata["bands"]):
 75        origin_raster_band = ref.GetRasterBand(band + 1)
 76        target_raster_band = shifted.GetRasterBand(band + 1)
 77
 78        target_raster_band.WriteArray(origin_raster_band.ReadAsArray())
 79
 80        if src_nodata is not None:
 81            if utils_base._check_variable_is_int(src_nodata):
 82                target_raster_band.SetNoDataValue(int(src_nodata))
 83            else:
 84                target_raster_band.SetNoDataValue(float(src_nodata))
 85
 86    shifted.FlushCache()
 87    shifted = None
 88    ref = None
 89
 90    return out_path
 91
 92
 93def raster_shift(
 94    raster: Union[str, gdal.Dataset, List],
 95    shift_list: List[Union[int, float]],
 96    out_path: Optional[str] = None,
 97    *,
 98    overwrite: bool = True,
 99    prefix: str = "",
100    suffix: str = "",
101    add_uuid: bool = False,
102    add_timestamp: bool = False,
103    creation_options: Optional[List[str]] = None,
104) -> Union[str, List[str], gdal.Dataset]:
105    """Shifts a raster in a given direction. (The frame is shifted)
106
107    Parameters
108    ----------
109    raster : Union[str, List, gdal.Dataset]
110        The raster(s) to be shifted.
111
112    shift_list : List[Union[int, float]]
113        The shift in x and y direction.
114
115    out_path : Optional[str], optional
116        The path to the output raster. If None, the raster is created in memory., default: None
117
118    overwrite : bool, optional
119        If True, the output raster will be overwritten if it already exists., default: True
120
121    prefix : str, optional
122        The prefix to be added to the output raster name., default: ""
123
124    suffix : str, optional
125        The suffix to be added to the output raster name., default: ""
126
127    add_uuid : bool, optional
128        If True, a unique identifier will be added to the output raster name., default: False
129
130    add_timestamp : bool, optional
131        If True, a timestamp will be added to the output raster name., default: False
132
133    creation_options : Optional[List[str]], optional
134        The creation options to be used when creating the output., default: None
135
136    Returns
137    -------
138    Union[str, List[str]]
139        The path(s) to the shifted raster(s).
140    """
141    utils_base._type_check(raster, [str, gdal.Dataset, [str, gdal.Dataset]], "raster")
142    utils_base._type_check(shift_list, [[int, float], tuple], "shift_list")
143    utils_base._type_check(out_path, [str, [str], None], "out_path")
144    utils_base._type_check(overwrite, [bool], "overwrite")
145    utils_base._type_check(creation_options, [[str], None], "creation_options")
146    utils_base._type_check(prefix, [str], "prefix")
147    utils_base._type_check(suffix, [str], "suffix")
148    utils_base._type_check(add_uuid, [bool], "add_uuid")
149    utils_base._type_check(add_timestamp, [bool], "add_timestamp")
150
151    input_is_list = isinstance(raster, list)
152
153    input_list = utils_io._get_input_paths(raster, "raster")
154    output_list = utils_io._get_output_paths(
155        input_list,
156        out_path,
157        overwrite=overwrite,
158        prefix=prefix,
159        suffix=suffix,
160        add_uuid=add_uuid,
161        add_timestamp=add_timestamp,
162        change_ext="tif",
163    )
164
165    utils_path._delete_if_required_list(output_list, overwrite)
166
167    shifted_rasters = []
168    for idx, in_raster in enumerate(input_list):
169        shifted_rasters.append(
170            _raster_shift(
171                in_raster,
172                shift_list,
173                out_path=output_list[idx],
174                overwrite=overwrite,
175                creation_options=creation_options,
176            )
177        )
178
179    if input_is_list:
180        return shifted_rasters
181
182    return shifted_rasters[0]
183
184
185def raster_shift_pixel(
186    raster: Union[str, gdal.Dataset],
187    shift_list: List[Union[int, float]],
188    out_path: Optional[str] = None,
189) -> str:
190    """Shifts a raster in a given direction. (The pixels are shifted, not the frame)
191
192    Parameters
193    ----------
194    raster : Union[str, gdal.Dataset]
195        The raster to be shifted.
196
197    shift_list : List[Union[int, float]]
198        The shift in x and y direction.
199
200    out_path : Optional[str], optional
201        The path to the output raster. If None, the raster is created in memory., default: None
202
203    Returns
204    -------
205    str
206        The path to the shifted raster.
207    """
208    utils_base._type_check(raster, [str, gdal.Dataset], "raster")
209    utils_base._type_check(shift_list, [[int, float], tuple], "shift_list")
210    utils_base._type_check(out_path, [str, None], "out_path")
211
212    arr = core_raster_io.raster_to_array(raster)
213
214    offsets, weights = kernel_shift(shift_list[0], shift_list[1])
215
216    arr_float32 = arr.astype(np.float32, copy=False)
217
218    for channel in range(arr.shape[2]):
219        arr[:, :, channel] = convolve_array_simple(
220            arr_float32[:, :, channel], offsets, weights,
221        )
222
223    if out_path is None:
224        out_path = utils_path._get_temp_filepath("shifted_raster.tif")
225    else:
226        if not utils_path._check_is_valid_output_filepath(out_path):
227            raise ValueError(f"out_path is not a valid output path: {out_path}")
228
229    utils_path._delete_if_required(out_path, overwrite=True)
230
231    return core_raster_io.array_to_raster(
232        arr,
233        reference=raster,
234        out_path=out_path,
235    )