def project_raster(ds, out_path, **kwargs):
ds, ras = ds_name(ds)
out_file = context_file(ras, out_path)
if os.path.exists(out_file):
return out_file
# input SpatialReference
in_srs = kwargs.pop('srcSRS', None)
inSpatialRef = read_srs([ds, in_srs])
# output SpatialReference
out_srs = kwargs.pop('dstSRS', "+proj=longlat +datum=WGS84 +ellps=WGS84")
outSpatialRef = read_srs(out_srs)
resample_alg = kwargs.pop('resampleAlg', gdal.GRA_Average)
option = gdal.WarpOptions(creationOptions=CREATION,
resampleAlg=resample_alg,
srcSRS=inSpatialRef,
dstSRS=outSpatialRef,
multithread=True,
**kwargs)
gdal.Warp(out_file, ds, options=option)
return out_file
def convert_uint8(ds, in_no_data=None, out_no_data=255):
ds, ras = ds_name(ds)
frist_band = ds.GetRasterBand(1)
if (frist_band.DataType != gdal.GDT_Byte
or frist_band.ReadAsArray(0, 0, 1, 1).dtype != np.int8):
return ras
if frist_band.GetNoDataValue() is not None:
in_no_data = frist_band.GetNoDataValue()
if in_no_data is None:
raise (ValueError("in_no_data must be initialed"))
option = gdal.WarpOptions(multithread=True,
creationOptions=CREATION,
outputType=gdal.GDT_Byte)
out_file = rep_file(os.path.dirname(ras), ras)
ds_out = gdal.Warp(out_file, ras, options=option)
for i in range(1, 1 + ds.RasterCount):
band = ds.GetRasterBand(i)
band_out = ds_out.GetRasterBand(i)
if band.GetNoDataValue() is None:
band.SetNoDataValue(in_no_data)
band_out.SetNoDataValue(out_no_data)
block_write(ds, band, band_out, _map_no_data)
ds = None
ds_out = None
gdal.GetDriverByName('GTiff').Delete(ras)
os.rename(out_file, ras)
return ras
def masked_outside(shp, ds):
try:
import psutil
except Exception:
print('psutil must be installed first')
ds, ras = ds_name(ds)
ds = gdal.Open(ras, gdal.GA_Update)
trans = ds.GetGeoTransform()
ratio = int(ds.RasterYSize * ds.RasterXSize /
(psutil.virtual_memory().available * 0.5 /
(2 + ds.ReadAsArray(0, 0, 1, 1).dtype.itemsize)) + 1)
# create the output layer
out_shp = '/vsimem/outline_wgs84.shp'
shp_projection(shp, out_shp, out_srs=read_srs(ds))
outDataSet = ogr.Open(out_shp)
outLayer = outDataSet.GetLayer()
band = ds.GetRasterBand(1)
b_xsize, b_ysize = band.GetBlockSize()
xsize = max(ds.RasterXSize // ratio + 1, b_xsize)
ysize = max(ds.RasterYSize // ratio + 1, b_ysize)
xoffs = range(0, ds.RasterXSize, xsize)
yoffs = range(0, ds.RasterYSize, ysize)
for i in range(ds.RasterCount):
band = ds.GetRasterBand(i + 1)
for xoff in xoffs:
for yoff in yoffs:
win_xsize = min(xsize, ds.RasterXSize - xoff)
win_ysize = min(ysize, ds.RasterYSize - yoff)
# Create the destination data source
target_ds = gdal.GetDriverByName('GTiff').Create(
'/vsimem/_outside.tif', win_xsize, win_ysize, 1, gdal.GDT_Byte)
left_up_lon, left_up_lat = imagexy2geo(ds, yoff - 0.5, xoff - 0.5)
target_ds.SetGeoTransform((left_up_lon, trans[1], 0, left_up_lat, 0, trans[5]))
target_ds.SetProjection(ds.GetProjection())
target_band = target_ds.GetRasterBand(1)
target_band.SetNoDataValue(0)
# Rasterize
gdal.RasterizeLayer(target_ds, [1], outLayer, burn_values=[1],
options=['ALL_TOUCHED=TRUE'])
outside = np.logical_not(target_ds.GetRasterBand(1).ReadAsArray())
target_band = None
target_ds = None
if outside.any():
arr = ds.ReadAsArray(xoff=xoff, yoff=yoff,
xsize=win_xsize, ysize=win_ysize)
arr[outside] = band.GetNoDataValue()
band.WriteArray(arr, xoff=xoff, yoff=yoff)
del arr
del outside
band = None
def resample(ds, out_path, **kwargs):
ds, ras = ds_name(ds)
out_file = context_file(ras, out_path)
if os.path.exists(out_file):
return out_file
resample_alg = kwargs.pop('resampleAlg', gdal.GRA_Average)
option = gdal.WarpOptions(multithread=True,
creationOptions=CREATION,
resampleAlg=resample_alg,
**kwargs)
gdal.Warp(out_file, ds, options=option)
return out_file
def grid_bound(ds, regions):
"""return gridded bound whose resolution exactly matches that of ds
Args:
ds (gdal.dataset): dataset or path
regions (list or str): regions of interest
Returns:
bound: spatial range of the regions in the ds
crs_transform: like bound but used in GEE
"""
ds = ds_name(ds)[0]
t = ds.GetGeoTransform()
bound, bound_srs = bound_layers(regions)
bound = bound_raster(ds, bound, bound_srs)[0]
crs_transform = [t[1], 0, bound[0], 0, t[5], bound[-1]]
return bound, crs_transform
def bound_raster(ds, bound, bound_srs="+proj=longlat +datum=WGS84 +ellps=WGS84"):
ds, _ = ds_name(ds)
t = ds.GetGeoTransform()
x_min, y_min, x_max, y_max = _prj_bound(ds, bound, bound_srs)
ulX, ulY = geo2imagexy(ds, x_min, y_min)
lrX, lrY = geo2imagexy(ds, x_max, y_max)
clip_range = [min(ulX, lrX), min(ulY, lrY),
abs(ulX - lrX) + 1, abs(ulY - lrY) + 1]
ul_lon = t[0] + t[1] * clip_range[0] + t[2] * clip_range[1]
ul_lat = t[3] + t[4] * clip_range[0] + t[5] * clip_range[1]
lr_lon = t[0] + t[1] * (clip_range[0] + clip_range[2]) + \
t[2] * (clip_range[1] + clip_range[3])
lr_lat = t[3] + t[4] * (clip_range[0] + clip_range[2]) + \
t[5] * (clip_range[1] + clip_range[3])
bound = [min(ul_lon, lr_lon), min(ul_lat, lr_lat),
max(ul_lon, lr_lon), max(ul_lat, lr_lat)]
return bound, read_srs(ds).ExportToProj4()
def mosaic(ras_paths, out_path, **kwargs):
ds = ras_paths[0]
ds, ras = ds_name(ds)
out_file = context_file(ras, out_path)
if os.path.exists(out_file):
return out_file
separate = kwargs.pop('separate', False)
resample_alg = kwargs.pop('resampleAlg', gdal.GRA_Average)
ds = gdal.BuildVRT('/vsimem/Mosaic.vrt', ras_paths, separate=separate)
option = gdal.WarpOptions(multithread=True,
creationOptions=CREATION,
resampleAlg=resample_alg,
**kwargs)
gdal.Warp(out_file, ds, options=option)
return out_file
def broadcast_args(ds_multi, calc_args, band_idxs):
iter_ds_multi, iter_calc_args, iter_band_idxs = check_iter(
ds_multi, calc_args, band_idxs)
if iter_ds_multi:
ds = ds_multi[0]
else:
ds = ds_multi
ds, ras = ds_name(ds)
if band_idxs is not None:
if iter_band_idxs and iter_calc_args:
if len(band_idxs) != len(calc_args):
raise Exception(
'length of band list not equal to that of calc args')
elif iter_band_idxs:
calc_args = [calc_args] * len(band_idxs)
elif iter_calc_args:
band_idxs = [band_idxs] * len(calc_args)
else:
calc_args = [calc_args]
band_idxs = [band_idxs]
else:
n_band = ds.RasterCount
if iter_calc_args:
if len(calc_args) != n_band:
raise Exception('calc args length not equal to band counts')
else:
calc_args = [calc_args] * n_band
if iter_ds_multi:
band_idxs = np.repeat(np.arange(1, n_band + 1, dtype=int),
len(ds_multi)).reshape(-1, len(ds_multi))
else:
band_idxs = np.arange(1, n_band + 1, dtype=int)
if not iter_ds_multi:
ds_multi = [ras]
band_idxs = np.array(band_idxs).reshape(len(band_idxs), 1)
return ds_multi, calc_args, band_idxs
def tif_copy_assign(out_file, ds_eg, array, srs=None, no_data=None):
if os.path.exists(out_file):
return out_file
ds_eg = ds_name(ds_eg)[0]
if array.ndim == 2:
array = array.reshape([1, *array.shape])
if array.ndim != 3:
raise (Exception('array must be 2 dims or 3 dims'))
# set nodata value
if no_data is None:
if ds_eg.GetRasterBand(1).GetNoDataValue() is not None:
no_data = ds_eg.GetRasterBand(1).GetNoDataValue()
else:
raise (Exception('nodata must be passed'))
if isinstance(array, np.ma.core.MaskedArray):
array.set_fill_value(no_data)
array = array.filled()
ds = gdal.GetDriverByName('GTiff').Create(out_file, array.shape[2],
array.shape[1], array.shape[0],
TYPE_MAP[array.dtype.name],
CREATION)
# fill with array
for i in range(1, 1 + ds.RasterCount):
band = ds.GetRasterBand(i)
band.SetNoDataValue(no_data)
band.WriteArray(array[i - 1])
# set geotransform
trans = ds_eg.GetGeoTransform()
ds.SetGeoTransform(tuple(trans))
# set SpatialReference
ds.SetSpatialRef(read_srs([srs, ds_eg]))
return out_file
def map_calc(ds_multi, calc_args, out_path, band_idxs=None, multiprocess=True):
iter_ds_multi = isinstance(ds_multi,
Iterable) and not isinstance(ds_multi, str)
if iter_ds_multi:
ds = ds_multi[0]
else:
ds = ds_multi
ds, ras = ds_name(ds)
out_file = context_file(ras, out_path)
if os.path.exists(out_file):
return out_file
ds_multi, calc_args, band_idxs = broadcast_args(ds_multi, calc_args,
band_idxs)
n = len(calc_args)
args = zip(np.arange(1, n + 1, dtype=int), [ds_multi] * n, band_idxs,
calc_args, [out_file] * n)
if multiprocess:
with Pool(min(cpu_count() - 1, n)) as p:
tem_files = p.starmap(band_map, args)
else:
tem_files = []
for arg in args:
tem_files.append(band_map(*arg))
if len(tem_files) == 1:
os.rename(tem_files[0], out_file)
else:
mosaic(tem_files, out_file, separate=True)
[os.remove(f) for f in tem_files]
return out_file
def grib_to_tif(ds, out_path=None, **kwargs):
ds, ras = ds_name(ds)
if os.path.splitext(os.path.basename(ras))[1] != '.grib':
return
if out_path:
out_file = context_file(ras, out_path)
else:
out_file = os.path.join(
os.path.dirname(ras),
os.path.splitext(os.path.basename(ras))[0] + '.tif')
if os.path.exists(out_file):
return out_file
srs = kwargs.pop('dstSRS', "+proj=longlat +datum=WGS84 +ellps=WGS84")
option = gdal.WarpOptions(multithread=True,
dstSRS=read_srs(srs),
creationOptions=CREATION,
**kwargs)
gdal.Warp(out_file, ds, options=option)
return out_file