如何利用带GPS坐标的.shp文件快速制作大.tif文件的二值掩码

问题背景 #

处理24GB的大型.tif文件，内存无法完全加载。尝试过rasterio官方的shp掩码方法，但批量处理所有多边形耗时超3小时，且常被服务器终止进程。现有代码仅能逐个处理多边形，效率低下。

当前代码核心问题 #

• 每次处理单个多边形都重复打开/写入输出TIFF，IO开销极大
• 重复调用rasterio.mask.mask和窗口读取，存在大量冗余计算
• 未利用rasterio分块处理能力，内存占用控制不佳

优化方案 #

优化方向1：按类别批量处理，减少IO操作 #

同一类别的所有多边形一次性处理，仅打开/写入输出TIFF一次，避免多次磁盘IO的耗时。

优化方向2：利用栅格化批量生成掩码 #

使用rasterio.features.rasterize直接批量栅格化同类别多边形，跳过冗余的影像读取步骤（若仅需掩码文件）。

优化方向3：压缩存储+坐标系预转换 #

• 掩码采用单通道+压缩存储，减小输出文件体积，加快读写速度
• 提前将矢量数据转换为栅格坐标系，避免每次处理时的动态转换开销

优化后代码 #

import rasterio
import geopandas as gpd
from tqdm import tqdm
import os
from rasterio.features import rasterize
from rasterio.windows import from_bounds, transform as window_transform

def generate_mask(raster_path, shape_path, output_path):
    # 确保输出目录存在
    os.makedirs(output_path, exist_ok=True)
    
    gdf = gpd.read_file(shape_path)
    with rasterio.open(raster_path) as src:
        # 统一坐标系，避免后续重复转换
        gdf = gdf.to_crs(src.crs)
        # 初始化掩码文件元数据：单通道、uint8类型、压缩存储
        mask_meta = src.meta.copy()
        mask_meta.update(count=1, dtype='uint8', compress='deflate')
        
        for class_num in tqdm(sorted(gdf['class'].unique())):
            output_file = os.path.join(output_path, f'{class_num}.tif')
            if os.path.exists(output_file):
                continue
            
            # 获取当前类别的所有多边形
            class_geoms = gdf[gdf['class'] == class_num]['geometry'].tolist()
            if not class_geoms:
                continue
            
            # 计算当前类别所有多边形的总边界，缩小处理范围
            total_bounds = gdf[gdf['class'] == class_num].total_bounds
            bbox = from_bounds(*total_bounds, src.transform)
            win_transform = window_transform(bbox, src.transform)
            win_h, win_w = int(bbox.height), int(bbox.width)
            
            # 批量栅格化生成掩码
            mask = rasterize(
                [(geom, 1) for geom in class_geoms],
                out_shape=(win_h, win_w),
                transform=win_transform,
                fill=0,
                dtype=rasterio.uint8
            )
            
            # 更新元数据并写入掩码文件
            mask_meta.update(height=win_h, width=win_w, transform=win_transform)
            with rasterio.open(output_file, 'w', **mask_meta) as dest:
                dest.write(mask, 1)
            print(f'类别 {class_num} 掩码已写入 {output_file}')

额外需求：结合原影像裁剪 #

如果需要生成原影像裁剪+掩码叠加的结果，可修改写入部分的代码：

import numpy as np

# 读取对应窗口的原影像
img_crop = src.read(window=bbox)
# 将单通道掩码扩展为与影像匹配的多通道
mask_3d = np.expand_dims(mask, axis=0)
# 应用掩码
out_image = img_crop * mask_3d

# 更新元数据为原影像的通道数和数据类型
img_meta = src.meta.copy()
img_meta.update(height=win_h, width=win_w, transform=win_transform, compress='deflate')
with rasterio.open(output_file, 'w', **img_meta) as dest:
    dest.write(out_image)

内容的提问来源于stack exchange，提问作者Nikita Artemenko