From ddd67a9206fc9f9ba4d1f8a032ecb26673a8caec Mon Sep 17 00:00:00 2001 From: Brendan Ward Date: Fri, 28 Mar 2014 15:25:15 -0700 Subject: [PATCH] Added support for rasterizing geometries (issue 45) and associated test cases. --- rasterio/_features.pyx | 77 ++++++++++++++++++++++- rasterio/_gdal.pxd | 5 +- rasterio/_ogr.pxd | 1 + rasterio/features.py | 58 ++++++++++++++++- rasterio/tests/test_features_rasterize.py | 51 +++++++++++++++ 5 files changed, 188 insertions(+), 4 deletions(-) create mode 100644 rasterio/tests/test_features_rasterize.py diff --git a/rasterio/_features.pyx b/rasterio/_features.pyx index 49585204..4ab2f137 100644 --- a/rasterio/_features.pyx +++ b/rasterio/_features.pyx @@ -1,7 +1,7 @@ # cython: profile=True import logging - +import json import numpy as np cimport numpy as np @@ -16,6 +16,9 @@ class NullHandler(logging.Handler): log.addHandler(NullHandler()) +ctypedef np.uint8_t DTYPE_UBYTE_t + + def _shapes(image, mask=None, connectivity=4, transform=None): """Return an iterator over Fiona-style features extracted from the image. @@ -198,7 +201,77 @@ def _sieve(image, size, connectivity=4, output=None): return out -ctypedef np.uint8_t DTYPE_UBYTE_t +def _rasterize_geometry_json(geometries, size_t rows, size_t columns, transform=None, all_touched=False, DTYPE_UBYTE_t default_value=1): + """ + :param geometries: array of either geometry json strings, or array of (geometry json string, value) pairs. + Values must be unsigned integer type. If not provided, this function will return a binary mask. + :param rows: number of rows + :param columns: number of columns + :param transform: GDAL style geotransform. If provided, will be set on output. + :param all_touched: if true, will rasterize all pixels touched, otherwise will use GDAL default method. + :param default_value: must be 8 bit unsigned, will be used to set all values not specifically passed in geometries. + """ + + cdef int retval + cdef size_t i + cdef size_t num_geometries = len(geometries) + cdef void *memdriver + cdef void *out_ds + cdef void *out_band + cdef void **ogr_geoms = NULL + cdef char **options = NULL + cdef double geotransform[6] + cdef double *pixel_values = NULL # requires one value per geometry + cdef int dst_bands[1] # only need one band to burn into + dst_bands[0] = 1 + + try: + if all_touched: + options = _gdal.CPLMalloc(sizeof(char*)) + options[0] = "ALL_TOUCHED=TRUE" + + #Do the boilerplate required to create a dataset, band, and set transformation + memdriver = _gdal.GDALGetDriverByName("MEM") + if memdriver == NULL: + raise ValueError("NULL driver for 'MEM'") + out_ds = _gdal.GDALCreate(memdriver, "output", columns, rows, 1, <_gdal.GDALDataType>1, NULL) #TODO: revisit data type + if out_ds == NULL: + raise ValueError("NULL output datasource") + if transform: + for i in range(6): + geotransform[i] = transform[i] + err = _gdal.GDALSetGeoTransform(out_ds, geotransform) + if err: + raise ValueError("transform not set: %s" % transform) + out_band = _gdal.GDALGetRasterBand(out_ds, 1) + if out_band == NULL: + raise ValueError("NULL output band") + + ogr_geoms = _gdal.CPLMalloc(num_geometries * sizeof(void*)) + pixel_values = _gdal.CPLMalloc(num_geometries * sizeof(double)) + for i in range(num_geometries): + entry = geometries[i] + if isinstance(entry, (tuple, list)): + geometry_json, value = entry + else: + geometry_json = entry + value = default_value #1 + ogr_geoms[i] = _ogr.OGR_G_CreateGeometryFromJson(geometry_json) + pixel_values[i] = value + + retval = _gdal.GDALRasterizeGeometries(out_ds, 1, dst_bands, num_geometries, ogr_geoms, NULL, geotransform, pixel_values, options, NULL, NULL) + out = np.zeros((rows, columns), np.uint8) + retval = io_ubyte(out_band, 0, 0, 0, columns, rows, out) + + finally: + _gdal.CPLFree(ogr_geoms) + _gdal.CPLFree(options) + _gdal.CPLFree(pixel_values) + if out_ds != NULL: + _gdal.GDALClose(out_ds) + + return out + cdef int io_ubyte( void *hband, diff --git a/rasterio/_gdal.pxd b/rasterio/_gdal.pxd index 4e39db9f..f84ea485 100644 --- a/rasterio/_gdal.pxd +++ b/rasterio/_gdal.pxd @@ -179,7 +179,10 @@ cdef extern from "gdal_alg.h": int GDALPolygonize(void *src_band, void *mask_band, void *layer, int fidx, char **options, void *progress_func, void *progress_data) int GDALSieveFilter(void *src_band, void *mask_band, void *dst_band, int size, int connectivity, char **options, void *progress_func, void *progress_data) - + int GDALRasterizeGeometries(void *out_ds, int band_count, int *dst_bands, int geom_count, void **geometries, + GDALTransformerFunc transform_func, void *transform, double *pixel_values, char **options, + void *progress_func, void *progress_data) + void *GDALCreateGenImgProjTransformer(void* hSrcDS, const char *pszSrcWKT, void* hDstDS, const char *pszDstWKT, int bGCPUseOK, double dfGCPErrorThreshold, diff --git a/rasterio/_ogr.pxd b/rasterio/_ogr.pxd index 2a5bd92c..9fb3796d 100644 --- a/rasterio/_ogr.pxd +++ b/rasterio/_ogr.pxd @@ -59,6 +59,7 @@ cdef extern from "ogr_api.h": void OGR_G_AddPoint_2D (void *geometry, double x, double y) void OGR_G_CloseRings (void *geometry) void * OGR_G_CreateGeometry (int wkbtypecode) + void * OGR_G_CreateGeometryFromJson(char *json) void OGR_G_DestroyGeometry (void *geometry) unsigned char * OGR_G_ExportToJson (void *geometry) void OGR_G_ExportToWkb (void *geometry, int endianness, char *buffer) diff --git a/rasterio/features.py b/rasterio/features.py index fa8e83fd..f642eb30 100644 --- a/rasterio/features.py +++ b/rasterio/features.py @@ -1,7 +1,8 @@ """Functions for working with features in a raster dataset.""" +import json import rasterio -from rasterio._features import _shapes, _sieve +from rasterio._features import _shapes, _sieve, _rasterize_geometry_json def shapes(image, mask=None, connectivity=4, transform=None): @@ -44,3 +45,58 @@ def sieve(image, size, connectivity=4, output=None): with rasterio.drivers(): return _sieve(image, size, connectivity) + +def rasterize_features(features, rows, columns, transform=None, all_touched=False, value_property=None): + """ + :param features: Fiona style features iterator (geojson python objects) + Values must be unsigned integer type. If not provided, this function will return a binary mask. + :param rows: number of rows + :param cols: number of columns + :param transform: GDAL style geotransform. If provided, will be set on output. + :param all_touched: if true, will rasterize all pixels touched, otherwise will use GDAL default method. + :param value_attribute: if provided, the name of the property to extract the values from for each feature + (must be unsigned integer type). If not provided, this function will return a binary mask. + """ + + + geoms = [] + for index, feature in enumerate(features): #have to loop over features, since it may be yielded from a generator + feature_json = json.dumps(feature['geometry']) + if value_property is not None: + if not not (feature['properties'] and value_property in feature['properties']): + raise ValueError("Value property is missing from feature number %i: %s" % (index, value_property)) + value = feature['properties'][value_property] + if not (isinstance(value, int) and value >= 0 and value < 256): + raise ValueError("Value for value_property is not valid for feature %i (must be 8 bit unsigned integer)" % index) + geoms.append((feature_json, value)) + else: + geoms.append(feature_json) + + with rasterio.drivers(): + return _rasterize_geometry_json(geoms, rows, columns, transform, all_touched) + + +def rasterize_geometries(geometries, rows, columns, transform=None, all_touched=False): + """ + :param geometries: array of Fiona style geometry objects or array of (geometry, value) pairs. + Values must be unsigned integer type. If not provided, this function will return a binary mask. + :param rows: number of rows + :param cols: number of columns + :param transform: GDAL style geotransform. If provided, will be set on output. + :param all_touched: if true, will rasterize all pixels touched, otherwise will use GDAL default method. + """ + + + geoms = [] + for index, entry in enumerate(geometries): #have to loop over features, since it may be yielded from a generator + if isinstance(entry, (tuple, list)): + geometry, value = entry + if not (isinstance(value, int) and value >= 0 and value < 256): + raise ValueError("Value for geometry number %i is not valid (must be 8 bit unsigned integer)" % index) + geoms.append((json.dumps(geometry), value)) + else: + geoms.append(json.dumps(entry)) + + with rasterio.drivers(): + return _rasterize_geometry_json(geoms, rows, columns, transform, all_touched) + diff --git a/rasterio/tests/test_features_rasterize.py b/rasterio/tests/test_features_rasterize.py new file mode 100644 index 00000000..f92694e6 --- /dev/null +++ b/rasterio/tests/test_features_rasterize.py @@ -0,0 +1,51 @@ +import logging +import sys +import numpy +import pytest +import rasterio +from rasterio.features import shapes, rasterize_geometries + +logging.basicConfig(stream=sys.stderr, level=logging.DEBUG) + + +def test_rasterize_geometries(): + rows = cols = 10 + transform = [0, 1, 0, 0, 0, 1] + geometry = {'type':'Polygon','coordinates':[[(2,2),(2,4.25),(4.25,4.25),(4.25,2),(2,2)]]} + + with rasterio.drivers(): + # we expect a subset of the pixels using default mode + result = rasterize_geometries([geometry], rows, cols, transform) + truth = numpy.zeros((rows, cols)) + truth[2:4, 2:4] = 1 + assert (result == truth).min() == True + + # we expect all touched pixels + result = rasterize_geometries([geometry], rows, cols, transform, all_touched=True) + truth = numpy.zeros((rows, cols)) + truth[2:5, 2:5] = 1 + assert (result == truth).min() == True + + # we expect the pixel value to match the one we pass in + value = 5 + result = rasterize_geometries([(geometry, value)], rows, cols, transform) + truth = numpy.zeros((rows, cols)) + truth[2:4, 2:4] = value + assert (result == truth).min() == True + + # we expect a ValueError if pixel value is not in 8 bit unsigned range + value = 500 + with pytest.raises(ValueError): + rasterize_geometries([(geometry, value)], rows, cols, transform) + + +def test_rasterize_geometries_symmetric(): + """Make sure that rasterize is symmetric with shapes""" + rows = cols = 10 + transform = [0, 1, 0, 0, 0, 1] + truth = numpy.zeros((rows, cols), dtype=rasterio.ubyte) + truth[2:5, 2:5] = 1 + with rasterio.drivers(): + s = shapes(truth, transform=transform) + result = rasterize_geometries(s, rows, cols, transform=transform) + assert (result == truth).min() == True