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 download_tiles(shp, tile_pixel):
# create the output layer
out_shp = '/vsimem/outline_wgs84.shp'
shp_projection(shp, out_shp)
outDataSet = ogr.Open(out_shp)
outLayer = outDataSet.GetLayer()
# Create the destination data source
x_res = int(360 / tile_pixel)
y_res = int(180 / tile_pixel)
target_ds = gdal.GetDriverByName('GTiff').Create(
'/vsimem/_tile.tif', x_res, y_res, 1, gdal.GDT_Byte)
target_ds.SetGeoTransform((-180, tile_pixel, 0, 90, 0, -tile_pixel))
target_ds.SetSpatialRef(read_srs('+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'))
band = target_ds.GetRasterBand(1)
band.SetNoDataValue(0)
# Rasterize
gdal.RasterizeLayer(target_ds, [1], outLayer, burn_values=[
1], options=['ALL_TOUCHED=TRUE'])
burn_tiles = target_ds.GetRasterBand(1).ReadAsArray()
rows, cols = np.where(burn_tiles == 1)
if rows.shape == 0:
return None
lon, lat = np.int_(imagexy2geo(target_ds, rows - 0.5, cols - 0.5))
return lon, lat
def polygonize(ds, shp_path):
if os.path.exists(shp_path):
return
# mapping between gdal type and ogr field type
type_mapping = {
gdal.GDT_Byte: ogr.OFTInteger,
gdal.GDT_UInt16: ogr.OFTInteger,
gdal.GDT_Int16: ogr.OFTInteger,
gdal.GDT_UInt32: ogr.OFTInteger,
gdal.GDT_Int32: ogr.OFTInteger,
gdal.GDT_Float32: ogr.OFTReal,
gdal.GDT_Float64: ogr.OFTReal,
gdal.GDT_CInt16: ogr.OFTInteger,
gdal.GDT_CInt32: ogr.OFTInteger,
gdal.GDT_CFloat32: ogr.OFTReal,
gdal.GDT_CFloat64: ogr.OFTReal
}
srcband = ds.GetRasterBand(1)
dst_layername = "Shape"
drv = ogr.GetDriverByName("ESRI Shapefile")
dst_ds = drv.CreateDataSource(shp_path)
dst_layer = dst_ds.CreateLayer(dst_layername, srs=read_srs(ds))
raster_field = ogr.FieldDefn('id', type_mapping[srcband.DataType])
dst_layer.CreateField(raster_field)
gdal.Polygonize(srcband, srcband, dst_layer, 0, [], callback=None)
return shp_path
def shp_geom_map(in_shp, out_shp, in_srs=None,
out_srs=None, idxs=None, func=None):
gdal.SetConfigOption("SHAPE_ENCODING", 'utf-8')
# Filename of input OGR file
driver = ogr.GetDriverByName("ESRI Shapefile")
if isinstance(in_shp, ogr.Layer):
inLayer = in_shp
else:
source_ds = ogr.Open(in_shp)
inLayer = source_ds.GetLayer()
# output SpatialReference
outSpatialRef = read_srs([out_srs, inLayer, in_srs])
# create the output layer
if 'vsimem' not in os.path.dirname(out_shp):
if (not os.path.exists(os.path.dirname(out_shp)) and
os.path.dirname(out_shp) != ''):
os.makedirs(os.path.dirname(out_shp))
if os.path.exists(out_shp):
driver.DeleteDataSource(out_shp)
outDataSet = driver.CreateDataSource(out_shp)
outLayer = outDataSet.CreateLayer(out_shp, outSpatialRef)
# add fields
inLayerDefn = inLayer.GetLayerDefn()
for i in range(0, inLayerDefn.GetFieldCount()):
fieldDefn = inLayerDefn.GetFieldDefn(i)
outLayer.CreateField(fieldDefn)
# get the output layer's feature definition
outLayerDefn = outLayer.GetLayerDefn()
# set all feature idxs if not assigned
if idxs is None:
idxs = range(inLayer.GetFeatureCount())
elif isinstance(idxs, int):
idxs = [idxs]
# loop through the input features
for i in idxs:
inFeature = inLayer.GetFeature(i)
# get the input geometry
geom = inFeature.GetGeometryRef()
# create a new feature
outFeature = ogr.Feature(outLayerDefn)
# set the geometry and attribute
outFeature.SetGeometry(geom if func is None else func(geom))
for i in range(0, outLayerDefn.GetFieldCount()):
outFeature.SetField(outLayerDefn.GetFieldDefn(
i).GetNameRef(), inFeature.GetField(i))
# add the feature to the shapefile
outLayer.CreateFeature(outFeature)
# dereference the features and get the next input feature
outFeature = None
return out_shp
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 bound_layers(layers):
# clip extent
if not isinstance(layers, list):
layers = [layers]
if isinstance(layers[0], str):
shp_in = [ogr.Open(lyr) for lyr in layers]
layers = [s.GetLayer() for s in shp_in]
win = np.array([lyr.GetExtent() for lyr in layers])
bound = [win[:, 0:2].min(), win[:, 2:].min(),
win[:, 0:2].max(), win[:, 2:].max()]
return bound, read_srs(layers).ExportToProj4()
def shp_projection(in_shp, out_shp, in_srs=None,
out_srs="+proj=longlat +datum=WGS84 +ellps=WGS84"):
gdal.SetConfigOption("SHAPE_ENCODING", 'utf-8')
# Filename of input OGR file
if isinstance(in_shp, ogr.Layer):
inLayer = in_shp
else:
source_ds = ogr.Open(in_shp)
inLayer = source_ds.GetLayer()
# input SpatialReference
inSpatialRef = read_srs([inLayer, in_srs])
# output SpatialReference
outSpatialRef = read_srs(out_srs)
# create the CoordinateTransformation
coordTrans = coord_trans(inSpatialRef, outSpatialRef)
return shp_geom_map(in_shp, out_shp,
out_srs=outSpatialRef,
func=lambda geom: geom.Transform(coordTrans))
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 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 zeros_tif(out_file,
x_size,
y_size,
n_band,
data_type,
trans,
srs,
no_data=2):
ds = gdal.GetDriverByName('GTiff').Create(out_file, x_size, y_size, n_band,
data_type, CREATION)
# fill with 0 and set no data
band = ds.GetRasterBand(1)
band.Fill(0)
band.SetNoDataValue(no_data)
# set geotransform
ds.SetGeoTransform(tuple(trans))
# set SpatialReference
ds.SetSpatialRef(read_srs(srs))
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
def extract(ras, shp, PROJ=None, no_data=None, stat=False, **kwargs):
if stat:
return _extract_stat(ras, shp, PROJ=PROJ, no_data=no_data, **kwargs)
if isinstance(ras, str):
ds = gdal.Open(ras)
else:
ds = ras
out_shp = '/vsimem/outline.shp'
shp_projection(shp, out_shp, out_srs=read_srs([ds, PROJ]))
outDataSet = ogr.Open(out_shp)
outLayer = outDataSet.GetLayer()
# set no data
if ds.GetRasterBand(1).GetNoDataValue() is not None:
no_data = ds.GetRasterBand(1).GetNoDataValue()
if no_data is None:
raise (ValueError("no_data must be initialed"))
# clip with shapefile
rect, _ = clip(ds, outLayer, no_data=no_data, **kwargs)
return rect.GetDescription()
def _extract_stat(ras, shp, PROJ=None, no_data=None, **kwargs):
if isinstance(ras, str):
ds = gdal.Open(ras)
else:
ds = ras
# Filename of input OGR file
driver = ogr.GetDriverByName("ESRI Shapefile")
source_ds = ogr.Open(shp)
inLayer = source_ds.GetLayer()
inLayerDefn = inLayer.GetLayerDefn()
# output SpatialReference
outSpatialRef = read_srs([ds, PROJ])
coordTrans = coord_trans(inLayer, outSpatialRef)
# initial output
row = inLayer.GetFeatureCount()
col = ds.RasterCount
stat = np.full((row, col), np.nan)
for r in range(row):
inFeature = inLayer.GetFeature(r)
# create the output layer
outFile = '/vsimem/outline.shp'
if os.path.exists(outFile):
driver.DeleteDataSource(outFile)
outDataSet = driver.CreateDataSource(outFile)
outLayer = outDataSet.CreateLayer(outFile, outSpatialRef)
# add fields
for i in range(0, inLayerDefn.GetFieldCount()):
fieldDefn = inLayerDefn.GetFieldDefn(i)
outLayer.CreateField(fieldDefn)
# get the output layer's feature definition
outLayerDefn = outLayer.GetLayerDefn()
# get the input geometry
geom = inFeature.GetGeometryRef()
# reproject the geometry
geom.Transform(coordTrans)
# create a new feature
outFeature = ogr.Feature(outLayerDefn)
# set the geometry and attribute
outFeature.SetGeometry(geom)
for i in range(0, outLayerDefn.GetFieldCount()):
outFeature.SetField(
outLayerDefn.GetFieldDefn(i).GetNameRef(),
inFeature.GetField(i))
# add the feature to the shapefile
outLayer.CreateFeature(outFeature)
if ds.GetRasterBand(1).GetNoDataValue() is not None:
no_data = ds.GetRasterBand(1).GetNoDataValue()
if no_data is None:
raise (ValueError("no_data must be initialed"))
rect = None
rect, burn_data = clip(ds, outLayer, no_data=no_data, **kwargs)
# iterate all bands
for c in range(1, col + 1):
rect_band = rect.GetRasterBand(c)
rect_data = rect_band.ReadAsArray()
no_data = rect_band.GetNoDataValue()
select = (rect_data != no_data)
try:
stat[r, c - 1] = np.average(rect_data[select],
weights=burn_data[select])
except ZeroDivisionError:
pass
return stat