代码实现:
注意针对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_buildings\train\Buildings'
ref_img_dir = r'Spacenet6_buildings\train\PS-RGB'
output_dir = r'Spacenet6_buildings\train\PS-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)