python分类前用样本生成mask

from osgeo import gdal, gdalconst
from osgeo import ogr
# 支持中文路径
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
rasterFile = r'G:\2020.tif'  # 原影像

shpFile = r'G:\all.shp'
dataset = gdal.Open(rasterFile, gdalconst.GA_ReadOnly)
geo_transform = dataset.GetGeoTransform()
#dataset.GetGeoTransform返回：左上角横坐标，水平空间分辨率，北方向旋转角度（无旋转为0），左上角纵坐标，同三，垂直方向分辨率

gdal.SetConfigOption("SHAPE_ENCODING", "UTF-8")
#打开矢量文件
shp = ogr.Open(shpFile, 0)
if shp is None:
    sys.exit('Could not open {0}.'.format(shpFile))
# layer_num = shp.GetLayerCount()#矢量文件中图层数
'''索引获取数据源中的图层方式'''
lyr = shp.GetLayer(0)
print(lyr.GetName())
m_layer = shp.GetLayerByIndex(0)#获取第一个图层，返回一个类对象
print(dir(lyr))
print(dir(m_layer))
fearture = lyr.GetFeature(1) #第二行属性的值
print(lyr.GetGeometryColumn())
cols = dataset.RasterXSize  # 列数
rows = dataset.RasterYSize  # 行数
x_min = geo_transform[0]
y_min = geo_transform[3]
pixel_width = geo_transform[1]
print(cols,rows)
# print(geo_transform)

target_ds = gdal.GetDriverByName('GTiff').Create(r"G:\mask.tif",
                                                 xsize=cols, ysize=rows, bands=1,
                                                 eType=gdal.GDT_Byte)
target_ds.SetGeoTransform(geo_transform)
target_ds.SetProjection(dataset.GetProjection())

# band = target_ds.GetRasterBand(1)
# band.SetNoDataValue(-999)
# band.FlushCache()
#矢量栅格化burn_values输出图像的像素值[0],options=["ATTRIBUTE=SHAPE_Leng"]选择矢量属性表中某个列名进行赋值
# gdal.RasterizeLayer(target_ds, [1], m_layer,burn_values=[1]) # 跟shp字段给栅格像元赋值,burn_values=[0]
gdal.RasterizeLayer(target_ds, [1], m_layer,options=["ATTRIBUTE= class_num"],burn_values=[0]) 

这个代码只是生成了和原始影像一毛一样的栅格mask代码。后面可以用arcgis把每类的shp转成栅格嵌套在这个mask上。