多波段影像数据区域统计（Python代码）

Tips ：本代码实现ArcGIS中的区域统计功能（Zonal），有两点特色：
1）能够将统计得到的属性值直接赋予矢量文件中；
2）可实现多波段影像区域统计（方法比较笨，先拆分波段，然后进行统计）。

区域统计的类型参考ArcGIS：
MEAN— Calculates the average of all cells in the value raster that belong to the same zone as the output cell.
MAJORITY — Determines the value that occurs most often of all cells in the value raster that belong to the same zone as the output cell.
MAX — Determines the largest value of all cells in the value raster that belong to the same zone as the output cell.
MEDIAN — Determines the median value of all cells in the value raster that belong to the same zone as the output cell.
MIN — Determines the smallest value of all cells in the value raster that belong to the same zone as the output cell.
MINORITY — Determines the value that occurs least often of all cells in the value raster that belong to the same zone as the output cell.
RANGE — Calculates the difference between the largest and smallest value of all cells in the value raster that belong to the same zone as the output cell.
STD — Calculates the standard deviation of all cells in the value raster that belong to the same zone as the output cell.
SUM — Calculates the total value of all cells in the value raster that belong to the same zone as the output cell.
VARIETY — Calculates the number of unique values for all cells in the value raster that belong to the same zone as the output cell. （unique）

# -*- coding: utf-8 -*-
"""
Created on Thu Dec 17 11:39:53 2020

@author: xiao_gf

"""
import ogr
from osgeo import gdal
import numpy as np
from rasterstats import zonal_stats


def readTif(fileName):
    dataset = gdal.Open(fileName)
    if dataset == None:
        print(fileName+"文件无法打开")
    return dataset

# 保存tif文件
def writeTiff(im_data, im_geotrans, im_proj, path):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    #创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if(dataset!= None):
        dataset.SetGeoTransform(im_geotrans) #写入仿射变换参数
        dataset.SetProjection(im_proj) #写入投影
    for i in range(im_bands):
        dataset.GetRasterBand(i+1).WriteArray(im_data[i])
    del dataset
   
def zonal(input_vector,in_raster,Stats_type,path):    
    in_ds =readTif(in_raster)
    im_width = in_ds.RasterXSize  # 栅格矩阵的列数
    im_height = in_ds.RasterYSize  # 栅格矩阵的行数
    im_bands = in_ds.RasterCount
    im_data = in_ds.ReadAsArray(0, 0, im_width, im_height)
    im_geotrans = in_ds.GetGeoTransform()
    im_proj = in_ds.GetProjection()
    
    for band in range(0,im_bands):
        data = im_data[band]
        rasters = path+'\B{}'.format(band+1)+'.tif'  #输出单波段名称为B1、B2、...,可根据自己需求修改设定
        writeTiff(data, im_geotrans, im_proj, rasters)
        
        shp = ogr.Open(input_vector,1)
        lyr = shp.GetLayer()   
        # 添加字段
        zonal_Field = ogr.FieldDefn("B{}".format(band+1), ogr.OFTReal) 
        #添加的字段名称和单波段文件名称一致，为B1、B2、...,可根据自己需求修改设定
        zonal_Field.SetPrecision(9)
        lyr.CreateField(zonal_Field)
            
        zonal_method = zonal_stats(input_vector, rasters, stats=[Stats_type])
        FID = 0
        for feat in lyr:
            Index = zonal_method[FID][Stats_type]
           # print(Index)
            feat.SetField("B{}".format(band+1), Index)
            feat.SetFID(FID)
            lyr.SetFeature(feat)
            FID +=1
        
if __name__ =="__main__":
    
    input_shp = r"D:\DATA\test.shp"  #输入统计的矢量文件
    input_raster = r"D:\DATA\test_r.tif"  #需要统计的栅格数据（遥感影像）
    path = r"D:\DATA\band" #获取的单波段数据存储地址
    Stats_type = 'mean' #区域统计的类型
    # Stats_type表示统计的类型，包括：'min', 'max', 'mean', 'count', 'sum', 'std', 'median', 'majority', 'minority', 'unique', 'range
    zonal(input_shp, input_raster, Stats_type, path)