Python 使用矢量数据对栅格进行裁剪

参考：GeospatialPython.com: Clip a Raster using a Shapefile

1. 主要步骤

• 将矢量的shapefile转化为mask数组；
• 加载栅格影像数据；
• 舍弃shapefile边界范围外的所有像元；
• 将shapefile范围外的所有像元值设置为NODATA;
• 可选步骤：对影像采用直方图拉伸，便于可视化
• 将裁剪后的影像保存为新的栅格数据

2. 工具

• GDAL:处理栅格数据
• Numpy：大型的多维数组运算
• PIL：便于影像与数组的相互转换

3. 代码

import operator

from cv2 import reduce
from osgeo import gdal, gdalnumeric, ogr, osr
import Image, ImageDraw


# 将栅格影像数据转化为数组
def imageToArray(img_file):
    arr = gdalnumeric.fromstring(img_file.tostring(), 'b')
    arr.shape = img_file.im.size[1], img_file.im.size[0]
    return arr


# 将数组转化为栅格影像
def arrayToImage(arr):
    img = Image.fromstring('L', (arr.shape[1], arr.shape[0]),
                         (arr.astype('b')).tostring())
    return img


# 计算像元的地理坐标信息
def world2Pixel(geoMatrix, x, y):
    ulX = geoMatrix[0]
    ulY = geoMatrix[3]
    xDist = geoMatrix[1]
    yDist = geoMatrix[5]
    rtnX = geoMatrix[2]
    rtnY = geoMatrix[4]
    pixel = int((x - ulX) / xDist)
    line = int((ulY - y) / yDist)
    return pixel, line


# 多维数组的直方图变换函数
def histogram(arr, bins=range(0, 256)):
    """
    :parameter arr: 输入的数组
    :parameter bins:直方图变换的输出范围
    """
    fa = arr.flat
    n = gdalnumeric.searchsorted(gdalnumeric.sort(fa), bins)
    n = gdalnumeric.concatenate([n, [len(fa)]])
    hist = n[1:] - n[:-1]
    return hist


# 直方图拉伸
def stretch(arr):
    hist = histogram(arr)
    im = arrayToImage(arr)
    lut = []
    for b in range(0, len(hist), 256):
        # step size
        step = reduce(operator.add, hist[b:b + 256]) / 255
        # create equalization lookup table
        n = 0
        for i in range(256):
            lut.append(n / step)
            n = n + hist[i + b]
    im = im.point(lut)
    return imageToArray(im)


if __name__ == "__main__":
    raster_file = "ras_test.tif"   # 待裁剪的栅格影像
    shp_file = "county"              # 用于裁剪的矢量数据
    output_file = "ras_clip"         # 裁剪输出后的数据

    srcArray = gdalnumeric.LoadFile(raster_file)    # 加载数据
    srcImage = gdal.Open(raster_file)               # 获取投影等地理信息
    geoTrans = srcImage.GetGeoTransform()

    shape_file = ogr.Open("%s.shp" % shp_file)
    lyr = shape_file.GetLayer(shp_file)
    poly = lyr.GetNextFeature()

    # 将layer的坐标信息转换为与栅格数据一致，并计算新影像的像元大小
    minX, maxX, minY, maxY = lyr.GetExtent()
    ulX, ulY = world2Pixel(geoTrans, minX, maxY)
    lrX, lrY = world2Pixel(geoTrans, maxX, minY)
    pxWidth = int(lrX - ulX)
    pxHeight = int(lrY - ulY)

    clip = srcArray[:, ulY:lrY, ulX:lrX]
    geoTrans = list(geoTrans)
    geoTrans[0] = minX
    geoTrans[3] = maxY

    # 制作mask数据
    points = []
    pixels = []
    geom = poly.GetGeometryRef()
    pts = geom.GetGeometryRef(0)
    for p in range(pts.GetPointCount()):
        points.append((pts.GetX(p), pts.GetY(p)))
    for p in points:
        pixels.append(world2Pixel(geoTrans, p[0], p[1]))
    rasterPoly = Image.new("L", (pxWidth, pxHeight), 1)
    rasterize = ImageDraw.Draw(rasterPoly)
    rasterize.polygon(pixels, 0)
    mask = imageToArray(rasterPoly)

    # 使用mask裁剪影像
    clip = gdalnumeric.choose(mask, (clip, 0)).astype(gdalnumeric.uint8)

    # 对3波段的数据进行拉伸
    for i in range(3):
        clip[i, :, :] = stretch(clip[i, :, :])

    # 将裁剪后的数据存为tif格式
    gdalnumeric.SaveArray(clip, "%s.tif" % output_file, format="GTiff", prototype=raster_file)

    # 将裁剪后的数据存为8bit的jpeg格式
    clip = clip.astype(gdalnumeric.uint8)
    gdalnumeric.SaveArray(clip, "%s.jpg" % output_file, format="JPEG")