GeoAI 大尺寸 GeoTIFF 平滑分块推理实战：predict_geotiff 与 D4 测试时增强完全指南

GeoAI 的 geoai.inference 模块为超大尺寸遥感影像（GeoTIFF）提供了一套内存高效、无拼接痕迹的分块推理（tiled inference）管线：通过重叠分块 + 平滑融合（blending）消除瓦片边界伪影，通过 windowed I/O 将内存占用控制在 batch_size × tile_size² 量级，并可选启用 D4 二面体群的 8 折测试时增强（TTA）提升预测质量。读完本文，你将掌握在任意 PyTorch 分割 / 回归模型上对大影像进行无缝推理的完整配置方法、四种融合策略的原理与选择、超分辨率推理接口，以及常见报错的规避技巧。

模块概览：解决什么问题

对高分辨率遥感影像做逐像素推理时，最直接的做法是把整幅影像一次性送入模型——但高分辨率影像往往有数亿像素，显存与内存都会瞬间爆掉；退而求其次的固定窗口切片拼接，则会在瓦片边界处产生肉眼可见的接缝伪影，因为模型在瓦片边缘的预测往往不如中心可靠。

geoai.inference 用四个机制同时解决这两个问题（对应 docs/inference.md 的核心说明）：

机制	作用
重叠分块（Overlapping tiles）	相邻瓦片共享一定像素宽度的边距，边界像素会被多个瓦片重复预测
平滑融合（Smooth blending）	2D 权重掩码让瓦片边缘的预测权重低、中心的权重高，加权平均后无缝拼接
窗口化 I/O（Windowed I/O）	通过 rasterio 的窗口（Window）按批读取瓦片，内存只与 batch_size × tile_size² 成正比，而非整幅影像
D4 测试时增强（TTA）	对每个瓦片施加 D4 二面体群的 8 个变换（恒等、3 次旋转、水平/垂直翻转、2 次对角翻转），推理后逆变换并求平均，以 8× 计算量换取预测精度提升

实现位于 geoai/inference.py，其公开 API 为：

• BlendMode——融合策略枚举
• create_weight_mask——生成 2D 融合权重掩码
• predict_geotiff——主推理入口（分割 / 回归通用）
• predict_geotiff_superres——超分辨率推理入口
• d4_forward / d4_inverse / d4_tta_forward——D4 变换与 TTA 工具函数

模块的公开符号由 geoai/init.py 通过 PEP 562 惰性导入机制暴露，因此既可以直接 from geoai.inference import ... 导入，也可以 from geoai import predict_geotiff 从包顶层导入；测试 tests/test_inference.py 专门验证了这两种导入方式的可用性。

安装与导入

pip install geoai-py

导入方式（两种等价）：

import numpy as np
import torch
import torch.nn as nn

# 方式一：从子模块导入
from geoai.inference import (
    BlendMode,
    create_weight_mask,
    predict_geotiff,
    d4_forward,
    d4_inverse,
    d4_tta_forward,
)

# 方式二：从包顶层惰性导入（依赖首次访问时才加载）
from geoai import predict_geotiff, BlendMode

四种融合策略与权重掩码

融合策略由 BlendMode 枚举定义，四种取值及其语义如下（源码见 geoai/inference.py）：

模式	权重语义	适用场景
NONE（"none"）	全像素权重恒为 1.0，重叠区直接平均、无锥形过渡	对比实验、验证接缝伪影来源
LINEAR（"linear"）	重叠区从边缘 0 到中心 1 的线性斜坡	通用、计算最轻的平滑方案
COSINE（"cosine"）	重叠区采用升余弦（Hann）锥形	过渡比 linear 更柔和
SPLINE（"spline"，默认）	升余弦的幂次（power 次方）锥形，过渡最平滑	默认推荐，接缝消除效果最佳

权重掩码的生成函数为：

def create_weight_mask(
    tile_size: int,
    overlap: int,
    mode: Union[str, BlendMode] = BlendMode.SPLINE,
    power: int = 2,
) -> np.ndarray

返回形状为 (tile_size, tile_size)、取值在 [0, 1] 的 float32 数组。掩码中心权重为 1、边缘趋近 0，相邻瓦片重叠区按权重做加权平均后，正好形成无缝过渡。

参数与约束（均有源码 / 测试依据）：

• overlap 必须满足 0 <= overlap < tile_size，否则抛出 ValueError（geoai/inference.py）；
• overlap = 0 或 mode = "none" 时返回全 1 掩码（等价于无平滑的普通平均）；
• spline 模式要求 overlap <= tile_size // 2，超出会明确提示改用 linear 或 cosine（geoai/inference.py）；
• linear / cosine 模式在 overlap > tile_size // 2 时仍能正确生成单调、对称、取值 <a href="https://link.gitcode.com/i/a0349ed46ece2060d3d2bfbdfcb128cf" target="_blank">0, 1] 的掩码——源码中用 np.minimum 规避了写入越界问题，[tests/test_inference.py 中的 test_large_overlap_no_corruption 对此有专门验证；
• 传入未知模式同样抛出 ValueError。

spline 模式内部由一维窗口 _spline_window_1d 生成（geoai/inference.py）：重叠区用升半余弦的 power 次方从 0 锥形上升到 1，非重叠中心区恒为 1。两个实现细节值得注意：endpoint=True 保证锥形端点精确到达 1.0，与平坦中心区之间不存在数值跳变；power 越大，边缘权重滚降越陡、权重越向中心集中。对应测试见 tests/test_inference.py。

create_weight_mask 也可以独立使用，用于自定义推理循环或可视化分析：

mask = create_weight_mask(256, 64, mode="spline")
assert mask.shape == (256, 256)

官方示例 docs/examples/smooth_inference.ipynb 中用 matplotlib 同时绘制了 none / linear / cosine / spline 四种掩码的热力图与中心行一维剖面，直观展示了重叠区（红色阴影带）内的锥形过渡差异。

predict_geotiff：主推理入口全参数详解

def predict_geotiff(
    model: "torch.nn.Module",
    input_raster: str,
    output_raster: str,
    tile_size: int = 256,
    overlap: int = 64,
    batch_size: int = 4,
    input_bands: Optional[List[int]] = None,
    num_classes: int = 1,
    output_dtype: str = "float32",
    output_nodata: float = -9999.0,
    blend_mode: Union[str, BlendMode] = "spline",
    blend_power: int = 2,
    tta: bool = False,
    preprocess_fn: Optional[Callable[..., np.ndarray]] = None,
    postprocess_fn: Optional[Callable[..., np.ndarray]] = None,
    device: Optional[str] = None,
    compress: str = "lzw",
    verbose: bool = True,
) -> str

完整签名与默认值见 geoai/inference.py，tests/test_inference.py 对全部参数和默认值做了签名级断言。各参数含义如下：

参数	默认值	说明与取值范围
model	必填	接受 (B, C, H, W) 浮点张量、返回 (B, num_classes, H, W) 的 PyTorch 模型
input_raster / output_raster	必填	输入 / 输出 GeoTIFF 路径；输入不存在抛 FileNotFoundError，输出目录不存在会自动创建
tile_size	256	方形瓦片边长（像素）
overlap	64	相邻瓦片重叠像素数，须满足 0 <= overlap < tile_size；越大越平滑但计算量越大
batch_size	4	每次前向传播的瓦片数；内存占用 ≈ batch_size × tile_size²
input_bands	None	1-based 波段索引列表；None 表示读取全部波段
num_classes	1	模型输出通道数；回归 / 二值分割用 1，多分类用实际类别数
output_dtype	"float32"	输出栅格 dtype，如 "float32"、"uint8"
output_nodata	-9999.0	输出 NoData 值；整数 dtype 时必须落在该 dtype 范围内
blend_mode	"spline"	融合策略："none" / "linear" / "cosine" / "spline"
blend_power	2	spline 融合的幂指数（其他模式忽略）
tta	False	是否启用 D4 测试时增强；开启后计算量约 8×
preprocess_fn	None	逐瓦片预处理回调 (np.ndarray) -> np.ndarray，输入形状 (C, H, W)；None 使用默认预处理
postprocess_fn	None	对最终融合结果 (num_classes, H, W) 的后处理回调（如分类 argmax）；None 为不处理
device	None	如 "cuda" / "cpu"；None 时按 CUDA > MPS > CPU 自动选择（见 geoai/utils/device.py）
compress	"lzw"	输出 GeoTIFF 压缩方式
verbose	True	打印进度信息

默认预处理逻辑（_default_preprocess，见 geoai/inference.py）：瓦片转 float32；若最大值 > 1.5 则除以 255 归一化到 [0, 1]；最后用 np.nan_to_num 把 NaN 置 0。如果你的影像不是 uint8（例如 uint16 的 NAIP 数据，取值范围到 10000），务必传入自定义 preprocess_fn。

底层执行流程

从源码 geoai/inference.py 可以还原出完整调用链：

1. 参数校验：输入文件存在性、overlap 范围、output_nodata 与 output_dtype 的兼容性；
2. 构建瓦片网格：stride = tile_size - overlap，从影像左上角按 stride 步进生成瓦片；边缘不足 tile_size 的瓦片会回拉起始坐标补满整瓦片，并用 dict.fromkeys 去重（边界处可能产生重复瓦片）；
3. 一次性生成权重掩码：create_weight_mask(tile_size, overlap, mode, power) 只算一次，供全部瓦片复用；
4. 按批窗口化读取：rasterio.windows.Window 只读取当前批的瓦片区域；边缘小瓦片零填充到 tile_size；逐瓦片调用 preprocess_fn；
5. 前向推理：torch.no_grad() 下调用 model(batch_tensor)，开启 tta 时改走 d4_tta_forward；
6. 首个批次形状校验：模型输出通道数必须等于 num_classes，否则抛出带指引的 ValueError；输出为 (B, H, W) 3D 形状时自动补为 (B, 1, H, W)；
7. 加权累加：预测 × 掩码后累加到 float64 的 output_sum 与 weight_sum 累加器；
8. 归一化：valid = weight_sum > 0 的像素用 output_sum / (weight_sum + 1e-8) 求加权平均，未覆盖像素（理论上不存在，但安全起见）写入 output_nodata；np.where 广播对 num_classes > 1 的多通道输出同样适用（测试 tests/test_inference.py 验证）；
9. 后处理与写出：可选 postprocess_fn，然后更新 profile（count / dtype / nodata / compress）并逐波段写盘。

实战：NAIP 影像无缝分割完整流程

下面流程完整来自官方示例 docs/examples/smooth_inference.ipynb，可直接照做。

1. 下载示例数据

import os
import geoai

url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train.tif"
input_path = "naip.tif"

if not os.path.exists(input_path):
    geoai.download_file(url, input_path)

geoai.print_raster_info(input_path)
geoai.view_raster(input_path)

示例使用一张小型 NAIP 航空影像（4 波段：RGB + NIR）进行演示。

2. 定义一个带全局上下文的演示模型

示例模型故意加入全局平均池化（GAP）——同一像素的预测会因所在瓦片不同而略有差异，这样在关闭融合时会产生肉眼可见的瓦片边界伪影，正好用于对比验证：

class SimpleSegModel(nn.Module):
    def __init__(self, in_channels: int = 4):
        super().__init__()
        self.features = nn.Conv2d(in_channels, 4, 3, padding=1, bias=True)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Conv2d(4, 1, 1, bias=True)
        self.sigmoid = nn.Sigmoid()
        # ...（权重按亮度、NDVI 等光谱特征手工固定）

    def forward(self, x):
        feat = self.features(x)
        feat = feat - self.gap(feat)   # 减去瓦片均值 -> 预测依赖瓦片
        return self.sigmoid(self.head(feat) * 5.0)

model = SimpleSegModel(in_channels=4)
model.eval()

3. 自定义预处理（关键）

NAIP 影像可能是 uint8（0–255）或 uint16（0–10000），默认预处理只处理 uint8，因此必须自定义：

def naip_preprocess(tile: np.ndarray) -> np.ndarray:
    """将 NAIP 瓦片 (C, H, W) 归一化到 float32 的 [0, 1]。"""
    tile = tile.astype(np.float32)
    max_val = tile.max()
    if max_val > 255:
        # uint16 NAIP（取值通常到 ~10000）
        tile = np.clip(tile, 0, 10000) / 10000.0
    elif max_val > 1.5:
        # uint8 NAIP（0-255）
        tile = tile / 255.0
    return np.nan_to_num(tile, nan=0.0)

4. 对比：无融合 vs spline 融合

# 无融合 —— 硬瓦片边界（后写覆盖）
predict_geotiff(
    model=model, input_raster=input_path, output_raster="output_none.tif",
    tile_size=256, overlap=64, batch_size=4, num_classes=1,
    blend_mode="none", preprocess_fn=naip_preprocess, verbose=True,
)

# spline 融合 —— 无缝输出
predict_geotiff(
    model=model, input_raster=input_path, output_raster="output_spline.tif",
    tile_size=256, overlap=64, batch_size=4, num_classes=1,
    blend_mode="spline", preprocess_fn=naip_preprocess, verbose=True,
)

用 rasterio 分别读取两张结果并用 matplotlib 并排显示，即可直观看到："none" 输出上瓦片接缝清晰可见，"spline" 输出完全无缝。

5. 启用 D4 测试时增强

predict_geotiff(
    model=model, input_raster=input_path, output_raster="output_tta.tif",
    tile_size=256, overlap=64, batch_size=4, num_classes=1,
    blend_mode="spline", tta=True,          # 启用 D4 TTA
    preprocess_fn=naip_preprocess, verbose=True,
)

6. 多类别分割输出

predict_geotiff 同样支持输出多通道（类别概率图）的模型：把 num_classes 设为模型输出通道数，并用 postprocess_fn 施加 argmax 得到类别图：

mc_model = MulticlassModel(in_channels=4, num_classes=5)  # 5 类模型，示例见 notebook
mc_model.eval()

predict_geotiff(
    model=mc_model,
    input_raster=input_path,
    output_raster="output_multiclass.tif",
    tile_size=256, overlap=64, batch_size=4,
    num_classes=5,
    output_dtype="float32",
    blend_mode="spline",
    preprocess_fn=naip_preprocess,
    postprocess_fn=lambda arr: np.argmax(arr, axis=0)[np.newaxis, :, :].astype(np.float32),
    verbose=True,
)

注意：postprocess_fn 的输入是融合完成后的 (num_classes, H, W) 数组，argmax 在融合后的概率图上执行，避免了对每块瓦片独立 argmax 再拼接所带来的边界类别跳变。

D4 测试时增强原理

D4 是正方形的二面体群，共 8 个对称变换：恒等、绕中心旋转 90°/180°/270°、水平翻转、垂直翻转、两条对角线的翻转。d4_forward 对一批 (B, C, H, W) 张量施加全部 8 个变换，d4_inverse 精确还原，d4_tta_forward 将二者组合成完整的"变换 → 推理 → 逆变换 → 平均"流程（实现见 geoai/inference.py）：

def d4_tta_forward(model, batch):
    augmented = d4_forward(batch)          # 8 个变换视图
    outputs = [model(aug) for aug in augmented]  # 8 次前向
    restored = d4_inverse(outputs)         # 逆变换回原始朝向
    return torch.stack(restored).mean(dim=0)     # 平均

三个函数也可作为独立工具用于自定义推理循环：

batch = torch.randn(1, 4, 64, 64)
augmented = d4_forward(batch)
print(len(augmented))            # 8

restored = d4_inverse(augmented)
for i, t in enumerate(restored):
    print(f"Transform {i}: max roundtrip error = {(t - batch).abs().max().item():.2e}")

geoai/inference.py 中最终输出是 8 个预测的平均；tests/test_inference.py 验证了：d4_forward 恒返回 8 个张量且形状不变、d4_inverse ∘ d4_forward 为精确恒等（roundtrip 误差为 0）、d4_tta_forward 恰好调用模型 8 次。代价是计算量增至 8×，因此 tta=True 适合"精度优先于速度"的场景。

超分辨率推理：predict_geotiff_superres

predict_geotiff_superres 是同一管线在超分辨率任务上的变体（geoai/inference.py）：模型输出空间分辨率高于输入（例如 ESRGAN 生成器），输入瓦片在低分辨率网格上读取，预测在高分辨率网格 (height×scale, width×scale) 上融合，仿射变换按 Affine.scale(1/scale, 1/scale) 同步缩小像素尺寸。

def predict_geotiff_superres(
    model, input_raster, output_raster, scale,
    tile_size=256, overlap=64, batch_size=4,
    input_bands=None, num_classes=1,
    output_dtype="float32", output_nodata=-9999.0,
    blend_mode="spline", blend_power=2, tta=False,
    preprocess_fn=None, postprocess_fn=None,
    device=None, compress="lzw", verbose=True,
) -> str

关键差异与约束：

• scale 必须 > 0；LR 侧的 tile_size / overlap / stride 会等比换算为 HR 侧的 hr_tile_size = tile_size * scale、hr_overlap = overlap * scale、hr_stride = stride * scale，权重掩码在 HR 网格上生成；
• 首个批次会校验模型输出空间尺寸必须等于 tile_size * scale，否则给出明确报错；
• 输出 transform 由 t = src.transform; hr_transform = t * rasterio.Affine.scale(1/scale, 1/scale) 更新，保持原点与旋转/剪切项一致。

在仓库中，该接口与 ESRGAN 生成器直接对接：geoai/esrgan.py 的 ESRGAN 类文档明确说明其 generator 可直接配合 geoai.inference.predict_geotiff_superres 使用（见 geoai/esrgan.py）。

边界情况与常见报错速查

以下错误均可在源码校验逻辑与 tests/test_inference.py 中找到对应依据：

场景	报错	解决方法
输入影像不存在	FileNotFoundError	检查 input_raster 路径
overlap >= tile_size 或 overlap < 0	ValueError	保证 0 <= overlap < tile_size（geoai/inference.py）
spline 模式且 overlap > tile_size // 2	ValueError	改用 linear / cosine，或缩小 overlap（geoai/inference.py）
output_dtype="uint8" 但 output_nodata=-9999.0	ValueError	整数 dtype 需用范围内的 nodata（如 0 或 255），见 tests/test_inference.py
模型输出通道数 ≠ num_classes	ValueError	把 num_classes 设为模型实际输出通道数
模型输出非 (B, H, W) / (B, C, H, W)	ValueError（superres 入口）	检查模型输出形状与 scale 是否匹配
未知 blend_mode	ValueError	仅支持 none / linear / cosine / spline

参数速查与调参建议

功能	参数	默认值	建议
瓦片大小	tile_size	256	按模型感受野与显存调整
重叠宽度	overlap	64	越大越平滑、计算量越大；spline 上限为 tile_size // 2
融合策略	blend_mode	"spline"	追求最平滑效果用 "spline"；大 overlap 用 "linear" / "cosine"
测试时增强	tta	False	精度优先于速度时置 True（8× 计算量）
自定义归一化	preprocess_fn	None（自动）	uint16 / 多光谱影像必须自定义
自定义后处理	postprocess_fn	None（不处理）	多分类分割用 argmax
输出 dtype	output_dtype	"float32"	类别图可改用 "uint8"（注意 nodata 范围）
设备	device	None（自动）	自动按 CUDA > MPS > CPU 选择

完整的可运行示例请直接查看 docs/examples/smooth_inference.ipynb，该 notebook 覆盖了从数据下载、掩码可视化、无融合/融合对比、D4 TTA 到多类别分割输出的全流程；模块级测试 tests/test_inference.py 则对签名、默认值、掩码性质、D4 变换与各类校验提供了可复现的验证。如果你需要基于已训练模型的影像级推理（而非本模块的分割 / 回归瓦片推理），仓库中 geoai/train.py 还提供了 inference_on_geotiff、instance_segmentation_inference_on_geotiff、multiclass_detection_inference_on_geotiff 等配套接口，可一并参考。