【rasterio】geojson矢量栅格化

使用rasterio将geojson矢量栅格化：

核心是使用rasterio.features.rasterize函数实现栅格化，具体考虑了：

1. 矢量化成单一值
2. 按照某个字段矢量化成不同的值
3. 对于空的矢量，矢量化得到0值的栅格
4. 写出栅格支持压缩

代码实现：

注意针对SN6数据做了特定修改

import os.path
from os.path import join, exists
import numpy as np
import rasterio as rio
from rasterio import features, enums
import geopandas as gpd
from tqdm import tqdm


def vector2img(vectorFileName, templateTifFileName, outputFileName, field=None):
    # Read in vector
    vector = gpd.read_file(vectorFileName)

    # Get list of geometries for all features in vector file
    geom = [shapes for shapes in vector.geometry]

    # Open example raster
    raster = rio.open(templateTifFileName)

    if len(geom) > 0:  # only rasterize non-empty vector
        if field is None:
            rasterized = features.rasterize(geom,
                                            out_shape=raster.shape,
                                            fill=0,
                                            out=None,
                                            transform=raster.transform,
                                            all_touched=False,
                                            default_value=1,
                                            dtype=None)
        else:
            # create a numeric unique value for each row
            vector[field] = range(0, len(vector))

            # create tuples of geometry, value pairs, where value is the attribute value you want to burn
            geom_value = ((geom, value) for geom, value in zip(vector.geometry, vector[field]))

            # Rasterize vector using the shape and transform of the raster
            rasterized = features.rasterize(geom_value,
                                            out_shape=raster.shape,
                                            transform=raster.transform,
                                            all_touched=True,
                                            fill=0,  # background value
                                            merge_alg=enums.MergeAlg.replace,
                                            dtype=np.int16)
    else:
        rasterized = np.zeros([raster.height, raster.width]).astype(np.uint8)

    with rio.open(
            outputFileName, "w",
            driver="GTiff",
            transform=raster.transform,
            dtype=rio.uint8,
            count=1,
            width=raster.width,
            height=raster.height,
            compress='lzw') as dst:
        dst.write(rasterized, indexes=1)


def vector2imgBatch(vector_dir, vector_postfix, ref_img_dir, img_postfix, output_dir, lab_postfix, field=None,
                    mismatch=False, vector_lastdir="", ref_img_lastdir=""):
    files = [f for f in os.listdir(vector_dir) if f.endswith(vector_postfix)]
    for file in tqdm(files):
        in_vector = join(vector_dir, file)
        if mismatch:
            ref_img = os.path.join(ref_img_dir, file[:-len(vector_postfix)].replace(vector_lastdir, ref_img_lastdir) + img_postfix)
        else:
            ref_img = os.path.join(ref_img_dir, file[:-len(vector_postfix)].replace(vector_lastdir, ref_img_lastdir) + img_postfix)
        out_label = os.path.join(output_dir,  file[:-len(vector_postfix)].replace(vector_lastdir, ref_img_lastdir) + lab_postfix)
        if exists(in_vector) and exists(ref_img):
            if os.path.exists(out_label):
                print('INFO: vector2img ' + out_label + " exists! Skip.")
            else:
                # print('vector2img: ' + file)
                vector2img(in_vector, ref_img, out_label, field)


if __name__ == '__main__':
    vector_dir = r'Spacenet6_buildingstrainBuildings'
    ref_img_dir = r'Spacenet6_buildingstrainPS-RGB'
    output_dir = r'Spacenet6_buildingstrainPS-RGB_Label'
    vector_postfix = 'geojson'
    img_postfix = 'tif'
    lab_postfix = 'tif'
    vector_lastdir = vector_dir.split('\')[-1]
    ref_img_lastdir = ref_img_dir.split('\')[-1]

    if not exists(output_dir):
        os.makedirs(output_dir)
    vector2imgBatch(vector_dir, vector_postfix, ref_img_dir, img_postfix, output_dir, lab_postfix,field=None,
                    mismatch=True, vector_lastdir=vector_lastdir, ref_img_lastdir=ref_img_lastdir)

参考：

Rasterize Vectors w. Rasterio — Python Open Source Spatial Programming & Remote Sensing (pygis.io)