使用rasterio的geometry_mask裁剪地图返回空图像，求代码修复

我尝试用geometry_mask结合分块与mmap控制内存占用实现地图裁剪，但当前代码返回空图像，循环未获取到任何数据，怀疑是block mask错误导致。预期是裁剪出多边形范围内的地图区域。

原代码：

import os
import numpy as np
import rasterio
from matplotlib import pyplot as plt
from rasterio.features import geometry_mask
from shapely.wkt import loads


def apply_mask_and_crop(polygon, map_image, transform_parameters, block_size=1024):
    # Load polygon from wkt strings
    polygon = loads(polygon)
    # Initialize boundaries
    ymin, ymax, xmin, xmax = float('inf'), 0, float('inf'), 0
    # Generate the transform
    west, south, east, north, map_width, map_height = transform_parameters
    transform = rasterio.transform.from_bounds(west, south, east, north, map_width, map_height)
    # Apply mask in blocks and find the bounding box of the mask
    for y in range(0, map_image.shape[0], block_size):
        for x in range(0, map_image.shape[1], block_size):
            y_end = min(y + block_size, map_image.shape[0])
            x_end = min(x + block_size, map_image.shape[1])
            # Generate the mask for the current block
            block_mask = geometry_mask([polygon], out_shape=(y_end - y, x_end - x), transform=transform, invert=True)
            # Apply the block mask directly to the map_image
            map_image[y:y_end, x:x_end][~block_mask] = 0
            # Update boundaries if there are any masked areas in the current block
            if np.any(block_mask):
                y_coords, x_coords = np.where(block_mask)
                ymin = min(ymin, y + y_coords.min())
                ymax = max(ymax, y + y_coords.max())
                xmin = min(xmin, x + x_coords.min())
                xmax = max(xmax, x + x_coords.max())
    # Ensure valid boundaries
    if ymin == float('inf'):
        ymin, xmin, ymax, xmax = 0, 0, 0, 0
    # Crop the image to the determined bounding box
    if ymin < ymax and xmin < xmax:
        cropped_image = map_image[ymin:ymax, xmin:xmax]
    else:
        cropped_image = np.empty((0, 0), dtype=map_image.dtype)
    # Calculate the geographic coordinates of the top-left and bottom-right
    top_left_geo = (ymin * transform.e + transform.f, xmin * transform.a + transform.c)
    bottom_right_geo = (ymax * transform.e + transform.f, xmax * transform.a + transform.c)
    return cropped_image, top_left_geo, bottom_right_geo


# Generate test data
multipolygon = "Multipolygon(((123 43,124 43,124 44,123 44,123 43)))"
image_shape, block_size = (24000, 24000, 3), 1024
image = np.memmap('temp', dtype=np.uint8, mode='w+', shape=image_shape)
for y in range(0, image_shape[0], block_size):
    for x in range(0, image_shape[1], block_size):
        y_end = min(y + block_size, image_shape[0])
        x_end = min(x + block_size, image_shape[1])
        block_data = np.random.randint(0, 256, size=(y_end - y, x_end - x, 3), dtype=np.uint8)
        image[y:y_end, x:x_end] = block_data
cropped_image, top_left_geo, bottom_right_geo = apply_mask_and_crop(multipolygon, image, (122.8, 42.5, 124.3, 44.2, 24000, 24000))
# Show image
plt.imshow(cropped_image)
plt.title("Image")
plt.axis('off')
plt.show()
image._mmap.close()
os.remove('temp')

核心错误原因 #

• block的transform参数错误：直接使用全局transform生成block mask，每个block对应的地理范围与全局范围不匹配，导致geometry_mask无法正确定位多边形在block中的位置。
• WKT格式不标准：输入的Multipolygon首字母小写，虽能被解析但存在兼容性风险。
• 三维图像mask适配缺失：原代码未处理RGB三通道图像的mask维度，导致mask无法正确应用。
• 地理坐标计算逻辑错误：手动计算地理坐标时误用transform参数，导致结果偏差。

修复后的完整代码 #

import os
import numpy as np
import rasterio
from matplotlib import pyplot as plt
from rasterio.features import geometry_mask
from shapely.wkt import loads


def apply_mask_and_crop(polygon, map_image, transform_parameters, block_size=1024):
    # 加载多边形，修正WKT格式兼容问题
    polygon = loads(polygon)
    # 初始化裁剪边界
    ymin, ymax, xmin, xmax = float('inf'), 0, float('inf'), 0
    # 生成全局transform
    west, south, east, north, map_width, map_height = transform_parameters
    global_transform = rasterio.transform.from_bounds(west, south, east, north, map_width, map_height)
    
    # 分块处理mask并更新裁剪边界
    for y in range(0, map_image.shape[0], block_size):
        for x in range(0, map_image.shape[1], block_size):
            y_end = min(y + block_size, map_image.shape[0])
            x_end = min(x + block_size, map_image.shape[1])
            # 计算当前block的左上角地理坐标
            block_x_geo, block_y_geo = rasterio.transform.xy(global_transform, y, x)
            # 生成block专属的transform
            block_transform = rasterio.transform.from_origin(
                block_x_geo, block_y_geo, 
                global_transform.a,  # x方向像素分辨率
                -global_transform.e  # y方向像素分辨率（取绝对值）
            )
            # 生成当前block的mask
            block_mask = geometry_mask(
                [polygon], 
                out_shape=(y_end - y, x_end - x), 
                transform=block_transform, 
                invert=True
            )
            # 将mask应用到RGB三通道图像
            map_image[y:y_end, x:x_end][~block_mask, :] = 0
            # 更新裁剪边界
            if np.any(block_mask):
                y_coords, x_coords = np.where(block_mask)
                ymin = min(ymin, y + y_coords.min())
                ymax = max(ymax, y + y_coords.max())
                xmin = min(xmin, x + x_coords.min())
                xmax = max(xmax, x + x_coords.max())
    # 处理无匹配区域的情况
    if ymin == float('inf'):
        ymin, xmin, ymax, xmax = 0, 0, 0, 0
    # 裁剪图像，+1避免丢失边缘像素
    if ymin < ymax and xmin < xmax:
        cropped_image = map_image[ymin:ymax+1, xmin:xmax+1].copy()
    else:
        cropped_image = np.empty((0, 0, 3), dtype=map_image.dtype)
    # 计算裁剪后图像的地理坐标
    top_left_geo = rasterio.transform.xy(global_transform, ymin, xmin)
    bottom_right_geo = rasterio.transform.xy(global_transform, ymax, xmax)
    return cropped_image, top_left_geo, bottom_right_geo


# 测试数据
# 修正WKT为标准大写格式
multipolygon = "MULTIPOLYGON(((123 43,124 43,124 44,123 44,123 43)))"
image_shape, block_size = (24000, 24000, 3), 1024
image = np.memmap('temp', dtype=np.uint8, mode='w+', shape=image_shape)
# 填充测试图像数据
for y in range(0, image_shape[0], block_size):
    for x in range(0, image_shape[1], block_size):
        y_end = min(y + block_size, image_shape[0])
        x_end = min(x + block_size, image_shape[1])
        block_data = np.random.randint(0, 256, size=(y_end - y, x_end - x, 3), dtype=np.uint8)
        image[y:y_end, x:x_end] = block_data
# 执行裁剪
cropped_image, top_left_geo, bottom_right_geo = apply_mask_and_crop(
    multipolygon, image, (122.8, 42.5, 124.3, 44.2, 24000, 24000)
)
# 显示结果
if cropped_image.size > 0:
    plt.imshow(cropped_image)
    plt.title("裁剪后地图")
    plt.axis('off')
    plt.show()
else:
    print("未找到匹配的裁剪区域")
# 清理临时文件
image._mmap.close()
os.remove('temp')

关键修复点说明 #

1. block专属transform生成：通过rasterio.transform.xy获取block左上角地理坐标，生成对应transform，确保mask与block位置匹配。
2. WKT格式修正：改为标准大写MULTIPOLYGON，避免解析兼容性问题。
3. 三维图像适配：给mask添加通道维度，确保能正确应用到RGB图像。
4. 裁剪边界优化：裁剪时给边界加1，避免丢失边缘像素。
5. 地理坐标计算优化：使用rasterio.transform.xy方法代替手动计算，避免参数理解错误。

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