Calculate raster statistics (min, max, mean) for each polygon intersecting a raster

raster_statistics.py

#!/usr/bin/env python

import numpy
try:
    from osgeo import gdal
except:
    import gdal
try:
    from osgeo import ogr
except:
    import ogr
try:
    from osgeo import gdalconst
except:
    import gdalconst
from PIL import Image, ImageDraw

raster = r'D:\MasterMap\terrain50_2013_07.tif'
band_number = 1

vector = r'D:\MasterMap\topo_area_su_region.shp'

dtype = numpy.float32

def get_raster_envelope(geo_transform, cols, rows):
    xmin = geo_transform[0]
    xmax = geo_transform[0] + cols*geo_transform[1] + rows*geo_transform[2]
    ymax = geo_transform[3]
    ymin = geo_transform[3] + cols*geo_transform[4] + rows*geo_transform[5]
    return xmin, xmax, ymin, ymax

def world2pixel(geo_transform, x, y):
    col = int((x - geo_transform[0]) / geo_transform[1])
    row = int((geo_transform[3] - y) / -geo_transform[5])
    return col, row
    
def bbox_contains(bbox1, bbox2):
    # does box 1 contain box 2?        
    if bbox1[0] <= bbox2[0] and bbox1[1] >= bbox2[1] and bbox1[2] <= bbox2[2] and bbox1[3] >= bbox2[3]:
        return True
    else:
        return False

# open the raster layer
raster_ds = gdal.Open(raster, gdalconst.GA_ReadOnly)
raster_band = raster_ds.GetRasterBand(band_number)
raster_geo_transform = raster_ds.GetGeoTransform()
raster_envelope = get_raster_envelope(raster_geo_transform, raster_ds.RasterXSize, raster_ds.RasterYSize)

# open the vector layer
vector_ds = ogr.Open(vector, gdalconst.GA_ReadOnly)
vector_layer = vector_ds.GetLayer()

results = []

feature_count = vector_layer.GetFeatureCount()

import time
t0 = time.time()

try:
    for n in xrange(0, feature_count):
        # clip raster to feature bbox
        feature = vector_layer.GetFeature(n)
        geometry = feature.GetGeometryRef()
        minX, maxX, minY, maxY = feature_envelope = geometry.GetEnvelope()
        if bbox_contains(raster_envelope, feature_envelope) is False:
            results.append(None)
            continue
        ulX, ulY = world2pixel(raster_geo_transform, minX, maxY)
        lrX, lrY = world2pixel(raster_geo_transform, maxX, minY)
        pxWidth = int(lrX - ulX)
        pxHeight = int(lrY - ulY)
        
        if pxWidth == 0 or pxHeight == 0:
            # feature is smaller than a single raster cell
            results.append(None)
            continue
                
        clip = raster_band.ReadAsArray(xoff=ulX, yoff=ulY, win_xsize=lrX-ulX, win_ysize=lrY-ulY, buf_obj=numpy.empty([pxHeight, pxWidth], dtype=dtype))
        
        clip_geo_transform = list(raster_geo_transform)
        clip_geo_transform[0] = minX
        clip_geo_transform[3] = maxY

        # create raster mask from feature
        points = geometry.GetGeometryRef(0)
        pixels = []
        for point in xrange(points.GetPointCount()):
            pixels.append(world2pixel(clip_geo_transform, points.GetX(point), points.GetY(point)))
        img = Image.new('L', (pxWidth, pxHeight), 0)
        ImageDraw.Draw(img).polygon(pixels, 1)
        mask = numpy.fromstring(img.tostring(), dtype=bool).reshape(clip.shape)
        
        # apply mask to clipped raster
        masked = clip[mask]
        
        if masked.size > 0:
            # calculate statistics                        
            result = masked.size, numpy.min(masked), numpy.max(masked), numpy.mean(masked)
        else:
            result = None
        
        #print result
except KeyboardInterrupt:
    pass

print n, time.time()-t0

raster_ds = None
vector_layer = None