批量裁剪栅格 python gdal

arcpy裁剪的缺点：
1.arcpy不支持中文
2.arcpy裁剪会对影像重采样，裁剪前的原图像元不对应，有偏移.

代码功能：
以一个shp裁剪多个栅格
输入：
1.shp全路径和为栅格命名的依据字段
2.栅格文件所在文件夹和栅格名称的后缀（.tif/.dat）
3.裁剪后的输出路径（文件夹）
输出：
在输出文件夹下以原栅格名称+‘_’+命名字段内容+‘_clip.tif’保存裁剪后的图像

import ogr
import os
import numpy
from tqdm import tqdm
import gdal
import time
from osgeo import gdalnumeric
from PIL import Image,ImageDraw

gdal.SetConfigOption("GDAL_ARRAY_OPEN_BY_FILENAME","TRUE")
# 栅格文件夹
rasterpath=r"******************"
# 栅格后缀名(.dat/.tif)
lastname=".dat"
#矢量文件
shp=r"**********.shp"
#命名字段(矢量文件的字段)
filename="Name"
#裁剪后文件存放位置（文件夹路径）
outputpath = r"*******************"

#将一个Python图像库的数组转换为一个gdal_array图片
def image2Array(i):
    a=gdalnumeric.fromstring(i.tobytes(),'b')
    a.shape=i.im.size[1],i.im.size[0]
    return a

# 数组写入dataset
def OpenArray(array,prototype_ds = None,xoff=0,yoff=0):
    ds = gdal.Open(gdalnumeric.GetArrayFilename(array))
    print(ds)
    if ds is not None and prototype_ds is not None:
        if type(prototype_ds).__name__ == 'str':
            prototype_ds = gdal.Open(prototype_ds)
        if prototype_ds is not None:
            gdalnumeric.CopyDatasetInfo(prototype_ds,ds,xoff=xoff,yoff=yoff)
    return ds

# 坐标换算
def world2Pixel(geoMatrix, x, y):
    ulx = geoMatrix[0]
    uly = geoMatrix[3]
    xDist = geoMatrix[1]
    yDist = geoMatrix[5]
    pixel = int((x - ulx) // xDist)
    line = int((uly - y) // abs(yDist))
    return (pixel, line)

strattime=time.time()

# 读取栅格
rasters = os.listdir(rasterpath)
rasterlist=list(filter(lambda x: x[-4:] == lastname, rasters))
print(rasterlist)
# 裁剪
with tqdm(total=len(rasterlist), iterable='iterable') as pbar:
    for ra in rasterlist:
        raster = str(rasterpath + '/' + ra)
        print(raster)

        # 将数据源作为gdal_array载入
        srcArray=gdalnumeric.LoadFile(raster)
        # 同时载入gdal库的图片从而获取geotransform
        srcImage=gdal.Open(raster)
        geoTrans=srcImage.GetGeoTransform()
        proj = srcImage.GetProjection()

        # 打开shp
        r = ogr.Open(shp)
        lyr = r.GetLayer()
        # 获取要素
        feature=lyr.GetNextFeature()
        while feature:
            geometry=feature.geometry()
            name=feature.GetField(filename)
            print(name)
            # 将图层扩展转换为图片像素坐标，需要每一个shp点的所在像素的左上角坐标
            minX,maxX,minY,maxY = geometry.GetEnvelope()

            # 计算要素四至对应的图片四至
            ulX, ulY = world2Pixel(geoTrans, minX, maxY)
            lrX, lrY = world2Pixel(geoTrans, maxX, minY)
            # 计算新图片的尺寸
            pxWidth = lrX - ulX
            pxHeight = lrY - ulY
            # 获取新图片影像数组
            clip = srcArray[ulY: lrY, ulX: lrX]
            print("clip.shape",clip.shape)
            # 计算新图片四至
            newOriginX = geoTrans[0]+ulX*geoTrans[1]
            newOriginY = geoTrans[3]+ulY*geoTrans[5]
            newEndX = geoTrans[0]+lrX*geoTrans[1]
            newEndY = geoTrans[3]+lrY*geoTrans[5]
            # 为新图片创建一个新的geomatrix对象以便附加地理参照数据
            newgeoTrans=list(geoTrans)
            newgeoTrans[0]=newOriginX
            newgeoTrans[3]=newOriginY

            # 使用PIL创建一个空白图片用于绘制多边形
            rasterPoly = Image.new('L', (pxWidth, pxHeight), 1)
            draw=ImageDraw.Draw(rasterPoly)
            # 在一个空白的8字节黑白掩膜图片上把点映射为像元绘制要素
            points = []
            pixels = []
            geom = feature.GetGeometryRef()
            cnt = geom.GetGeometryCount()

            if cnt > 1:#一个要素有多个元素(例如中国shp上的大陆与海南岛)
                for n in range(0,cnt):   
                    pos = geom.GetGeometryRef(n)#依次获取要素的部分元素polygon
                    pts = pos.GetGeometryRef(0)# 获取linestring
                    print("Geometry:", n, "points:", pts.GetPointCount())
                    for i in range(pts.GetPointCount()):#获取point
                        points.append((pts.GetX(i), pts.GetY(i)))
    #                 print(len(points))
                    for p in points:
                        p1,p2 = world2Pixel(newgeoTrans, p[0], p[1])
                        pixels.append((p1,p2)) 
    #                 print(len(pixels))
                    # 用像元位置绘制多边形 
                    draw.polygon(pixels,fill=0)
                    print("Geometry", n, "PILsize:",draw)
                    points = []
                    pixels = []

            else:  # 要素只有一个元素
                pts = geom.GetGeometryRef(0)
                # 获取所有点
                for i in range(pts.GetPointCount()):
                    points.append((pts.GetX(i), pts.GetY(i)))
                # 点坐标转图片坐标
                for p in points:
                    p1, p2 = world2Pixel(newgeoTrans, p[0], p[1])
                    pixels.append((p1, p2))
                # print(len(pixels))
                # 用像元位置绘制多边形          
                draw.polygon(pixels, fill=0)
                print("PILsize:", draw)

            # 使用PIL图片转换为Numpy掩膜数组
            mask = numpy.array(rasterPoly)
            # 根据掩膜图层对图像进行裁剪
            clip = gdalnumeric.numpy.choose(mask,(clip,0)).astype(gdalnumeric.numpy.float32)

            # 保存裁剪后的成果
            gtiffDriver = gdal.GetDriverByName('GTiff')
            if gtiffDriver is None:
                raise ValueError("Can't found Geotiff Driver")
            output = outputpath+(raster.split("/")[-1]).split(lastname)[0]+'_'+name+"_clip.tiff"
            # 数组写入dataset
            out = OpenArray(clip, prototype_ds=raster, xoff=ulX, yoff=ulY)
            print("out:", out)
            # 保存dataset
            outds = gtiffDriver.CreateCopy(output, out)
            # 写入坐标
            outds.SetGeoTransform(newgeoTrans)
            # 设置投影
            outds.SetProjection(proj)
            del out, outds, geometry, clip, mask
            feature = lyr.GetNextFeature()

        del feature, srcArray, srcImage, geoTrans, proj, r, lyr
        pbar.update(1)

endtime = time.time()
print("spend:", endtime-strattime)