编程_矢量转为栅格

ytkz2023-12-262024-09-09

需求

现在有一份SHP文件，需要把这份矢量文件制作为深度学习的标签。这时候第一个步骤是：把矢量栅格化。

代码

#!/usr/bin/env python
# -*- coding: utf-8 -*- 
# @Time : 2023/11/30 10:53 
# @File : shp2tif.py

'''
shp to tif
'''
from osgeo import gdal, ogr

def shp_to_tiff(shp_file,refore_tif, output_tiff):
    """
     将shp文件转换成tiff文件
    :param shp_file: 边界矢量数据
    :param output_tiff: 转换的tiff结果
    :param projection:  转换后的tiff格式
    :return: 矢量栅格图像
    """
    # 读取shp文件
    driver = ogr.GetDriverByName("ESRI Shapefile")
    data_source = driver.Open(shp_file, 1)
    # 获取图层文件对象
    shp_layer = data_source.GetLayer()

    img = gdal.Open(refore_tif)
    projection = img.GetProjection()
    transform = img.GetGeoTransform()
    cols = img.RasterXSize
    rows = img.RasterYSize

    # 根据模板tif属性信息创建对应标准的目标栅格
    target_ds = gdal.GetDriverByName('GTiff').Create(output_tiff, cols, rows, 1, gdal.GDT_Byte)
    target_ds.SetGeoTransform(transform)
    target_ds.SetProjection(projection)

    band = target_ds.GetRasterBand(1)
    # 设置背景数值
    NoData_value = 0
    band.SetNoDataValue(NoData_value)
    band.FlushCache()

    # 调用栅格化函数。gdal.RasterizeLayer函数有四个参数，分别有栅格对象，波段，矢量对象，value的属性值将为栅格值
    gdal.RasterizeLayer(target_ds, [1], shp_layer, options=["ATTRIBUTE=ID"])
    # gdal.RasterizeLayer(target_ds, [1], shp_layer)

    y_buffer = band.ReadAsArray()
    target_ds.WriteRaster(0, 0, cols, rows, y_buffer.tobytes())
    target_ds = None  # todo 释放内存，只有强制为None才可以释放干净
    del target_ds, shp_layer

if __name__ == '__main__':

    shp_file = r'D:\植被范围.shp'
    reference_tif = r'D:\img1.tif'
    output_tiff = r'D:\植被范围3.tif'
    shp_to_tiff(shp_file, reference_tif, output_tiff)

代码解析

1.读取shp文件
2.获取shp文件的图层对象
3.读取tif文件
4.获取tif文件的投影信息、仿射变换信息、行列数
5.根据tif文件的属性信息创建对应标准的目标栅格
6.设置背景数值
7.调用栅格化函数。gdal.RasterizeLayer函数有四个参数，分别有栅格对象，波段，矢量对象，value的属性值将为栅格值
8.将栅格数据写入到目标文件中
9.释放内存，只有强制为None才可以释放干净

结果如下：

矢量中的数值为小数

把shp_to_tiff函数增加多一个形参type，类型是字符串。

如果type=’int’，则表示矢量的数值是整数。

如果type=’float’，则表示矢量的数值是小数。

if type == 'int':
    target_ds = gdal.GetDriverByName('GTiff').Create(output_tiff, cols, rows, 1, gdal.GDT_Byte)
elif type == 'float':
    target_ds = gdal.GetDriverByName('GTiff').Create(output_tiff, cols, rows, 1, gdal.GDT_Float32)

还有，别忘了代码里的属性名字要与shp的数值那一列的名字要对应起来。

完整代码

#!/usr/bin/env python
# -*- coding: utf-8 -*- 
# @Time : 2023/12/26 22:53 
# @File : shp2tif.py

'''
shp to tif
'''
from osgeo import gdal, ogr


def shp_to_tiff(shp_file, refore_tif, output_tiff, type='int'):
    """
     将shp文件转换成tiff文件
    :param shp_file: 边界矢量数据
    :param output_tiff: 转换的tiff结果
    :param projection:  转换后的tiff格式
    :return: 矢量栅格图像
    """
    # 读取shp文件
    driver = ogr.GetDriverByName("ESRI Shapefile")
    data_source = driver.Open(shp_file, 1)
    # 获取图层文件对象
    shp_layer = data_source.GetLayer()

    img = gdal.Open(refore_tif)
    projection = img.GetProjection()
    transform = img.GetGeoTransform()
    cols = img.RasterXSize
    rows = img.RasterYSize

    # 根据模板tif属性信息创建对应标准的目标栅格
    if type == 'int':
        target_ds = gdal.GetDriverByName('GTiff').Create(output_tiff, cols, rows, 1, gdal.GDT_Byte)
    elif type == 'float':
        target_ds = gdal.GetDriverByName('GTiff').Create(output_tiff, cols, rows, 1, gdal.GDT_Float32)


    target_ds.SetGeoTransform(transform)
    target_ds.SetProjection(projection)

    band = target_ds.GetRasterBand(1)
    # 设置背景数值
    NoData_value = 0
    band.SetNoDataValue(NoData_value)
    band.FlushCache()

    # 调用栅格化函数。gdal.RasterizeLayer函数有四个参数，分别有栅格对象，波段，矢量对象，value的属性值将为栅格值
    gdal.RasterizeLayer(target_ds, [1], shp_layer, options=["ATTRIBUTE=ID"])

    # gdal.RasterizeLayer(target_ds, [1], shp_layer)

    y_buffer = band.ReadAsArray()
    target_ds.WriteRaster(0, 0, cols, rows, y_buffer.tobytes())
    target_ds = None  # todo 释放内存，只有强制为None才可以释放干净
    del target_ds, shp_layer


if __name__ == '__main__':
    shp_file = r'D:\Examplet_float.shp'
    reference_tif = r'D:\img1.tif'
    output_tiff = r'D:\raster5.tif'
    shp_to_tiff(shp_file, reference_tif, output_tiff, type='float')