Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Save xarray to a Cloud Optimized GeoTIFF (COG)
import gdal
from osgeo import gdal_array
import osr
import numpy
def get_dst_dataset(dst_img, cols, rows, layers, dtype, proj, gt):
"""
Create a GDAL data set in Cloud Optimized GeoTIFF (COG) format
:param dst_img: Output filenane full path
:param cols: Number of columns
:param rows: Number of rows
:param layers: Number of layers
:param dtype: GDAL type code
:param proj: Projection information in WKT format
:param gt: GeoTransform tupple
:return dst_ds: GDAL destination dataset object
"""
gdal.UseExceptions()
try:
# Default driver options to create a COG
driver = gdal.GetDriverByName('GTiff')
driver_options = ['COMPRESS=DEFLATE',
'BIGTIFF=YES',
'PREDICTOR=1',
'TILED=YES',
'COPY_SRC_OVERVIEWS=YES']
# Create driver
dst_ds = driver.Create(dst_img, cols, rows, layers,
dtype, driver_options)
# Set cartographic projection
dst_ds.SetProjection(proj)
dst_ds.SetGeoTransform(gt)
except Exception as err:
if err.err_level >= gdal.CE_Warning:
print('Cannot write dataset: %s' % self.input.value)
# Stop using GDAL exceptions
gdal.DontUseExceptions()
raise RuntimeError(err.err_level, err.err_no, err.err_msg)
gdal.DontUseExceptions()
return dst_ds
def get_resolution_from_xarray(xarray):
"""
Method to create a tuple (x resolution, y resolution) in x and y
dimensions.
:param xarray: xarray with latitude and longitude variables.
:return: tuple with x and y resolutions
"""
x_res = xarray.longitude.values[1] - xarray.longitude.values[0]
y_res = xarray.latitude.values[1] - xarray.latitude.values[0]
return (x_res, y_res)
def save_xarray(fname, xarray, data_var):
"""
Saves an xarray dataset to a Cloud Optimized GeoTIFF (COG)
:param fname: Full path of file where to save the data
:param xarray: xarray Dataset
:param data_var: Data variable in the xarray dataset
"""
# Create GeoTransform - perhaps the user requested a
# spatial subset, therefore is compulsory to update it
# GeoTransform -- case of a "north up" image without
# any rotation or shearing
# GeoTransform[0] top left x
# GeoTransform[1] w-e pixel resolution
# GeoTransform[2] 0
# GeoTransform[3] top left y
# GeoTransform[4] 0
# GeoTransform[5] n-s pixel resolution (negative value)
# Create tmp xarray DataArray
_xarray = getattr(xarray, data_var)
x_res, y_res = get_resolution_from_xarray(_xarray)
gt = (_xarray.longitude.data[0] - (x_res / 2.),
x_res, 0.0,
_xarray.latitude.data[0] - (y_res / 2.),
0.0, y_res)
# Coordinate Reference System (CRS) in a PROJ4 string to a
# Spatial Reference System Well known Text (WKT)
crs = xarray.attrs['crs']
srs = osr.SpatialReference()
srs.ImportFromProj4(crs)
proj = srs.ExportToWkt()
# Get GDAL datatype from NumPy datatype
if _xarray.dtype == 'bool':
dtype = gdal.GDT_Byte
else:
dtype = gdal_array.NumericTypeCodeToGDALTypeCode(_xarray.dtype)
# Dimensions
layers, rows, cols = _xarray.shape
# Create destination dataset
dst_ds = get_dst_dataset(dst_img=fname, cols=cols, rows=rows,
layers=layers, dtype=dtype, proj=proj, gt=gt)
for layer in range(layers):
dst_band = dst_ds.GetRasterBand(layer + 1)
# Date
if 'time' in _xarray.dims:
dst_band.SetMetadataItem('time',
_xarray.time.data[layer].astype(str))
# Data variable name
dst_band.SetMetadataItem('data_var', data_var)
# Data
dst_band.WriteArray(_xarray[layer].data)
dst_ds = None
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment