def write_output(raster, output, image_ds, gdal_frmt, ndv, band_names=None):
""" Write raster to output file """
from osgeo import gdal, gdal_array
logger.debug('Writing output to disk')
driver = gdal.GetDriverByName(str(gdal_frmt))
if len(raster.shape) > 2:
nband = raster.shape[2]
else:
nband = 1
ds = driver.Create(
output,
image_ds.RasterXSize, image_ds.RasterYSize, nband,
gdal_array.NumericTypeCodeToGDALTypeCode(raster.dtype.type)
)
if band_names is not None:
if len(band_names) != nband:
logger.error('Did not get enough names for all bands')
sys.exit(1)
if raster.ndim > 2:
for b in range(nband):
logger.debug(' writing band {b}'.format(b=b + 1))
ds.GetRasterBand(b + 1).WriteArray(raster[:, :, b])
ds.GetRasterBand(b + 1).SetNoDataValue(ndv)
if band_names is not None:
ds.GetRasterBand(b + 1).SetDescription(band_names[b])
ds.GetRasterBand(b + 1).SetMetadata({
'band_{i}'.format(i=b + 1): band_names[b]
})
else:
logger.debug(' writing band')
ds.GetRasterBand(1).WriteArray(raster)
ds.GetRasterBand(1).SetNoDataValue(ndv)
if band_names is not None:
ds.GetRasterBand(1).SetDescription(band_names[0])
ds.GetRasterBand(1).SetMetadata({'band_1': band_names[0]})
ds.SetProjection(image_ds.GetProjection())
ds.SetGeoTransform(image_ds.GetGeoTransform())
ds = None
def write_output(raster, output, image_ds, gdal_frmt, ndv, band_names=None):
""" Write raster to output file """
from osgeo import gdal, gdal_array
logger.debug('Writing output to disk')
driver = gdal.GetDriverByName(str(gdal_frmt))
if len(raster.shape) > 2:
nband = raster.shape[2]
else:
nband = 1
ds = driver.Create(
output, image_ds.RasterXSize, image_ds.RasterYSize, nband,
gdal_array.NumericTypeCodeToGDALTypeCode(raster.dtype.type))
if band_names is not None:
if len(band_names) != nband:
logger.error('Did not get enough names for all bands')
sys.exit(1)
if raster.ndim > 2:
for b in range(nband):
logger.debug(' writing band {b}'.format(b=b + 1))
ds.GetRasterBand(b + 1).WriteArray(raster[:, :, b])
ds.GetRasterBand(b + 1).SetNoDataValue(ndv)
if band_names is not None:
ds.GetRasterBand(b + 1).SetDescription(band_names[b])
ds.GetRasterBand(b + 1).SetMetadata(
{'band_{i}'.format(i=b + 1): band_names[b]})
else:
logger.debug(' writing band')
ds.GetRasterBand(1).WriteArray(raster)
ds.GetRasterBand(1).SetNoDataValue(ndv)
if band_names is not None:
ds.GetRasterBand(1).SetDescription(band_names[0])
ds.GetRasterBand(1).SetMetadata({'band_1': band_names[0]})
ds.SetProjection(image_ds.GetProjection())
ds.SetGeoTransform(image_ds.GetGeoTransform())
ds = None
def get_image_attribute(image_filename):
""" Use GDAL to open image and return some attributes
Args:
image_filename (string): image filename
Returns:
tuple (int, int, int, type): nrow, ncol, nband, NumPy datatype
"""
try:
image_ds = gdal.Open(image_filename, gdal.GA_ReadOnly)
except:
logger.error("Could not open example image dataset ({f})".format(f=image_filename))
sys.exit(1)
nrow = image_ds.RasterYSize
ncol = image_ds.RasterXSize
nband = image_ds.RasterCount
dtype = gdal_array.GDALTypeCodeToNumericTypeCode(image_ds.GetRasterBand(1).DataType)
return (nrow, ncol, nband, dtype)
def find_results(location, pattern):
""" Create list of result files and return sorted
Args:
location (str): directory location to search
pattern (str): glob style search pattern for results
Returns:
results (list): list of file paths for results found
"""
# Note: already checked for location existence in main()
records = []
for root, dirnames, filenames in os.walk(location):
for filename in fnmatch.filter(filenames, pattern):
records.append(os.path.join(root, filename))
if len(records) == 0:
logger.error('Error: could not find results in: {0}'.format(location))
sys.exit(1)
records.sort()
return records
