如何正确保存RasterVision语义分割预测结果？预测异常排查

问题背景 #

我已使用Rastervision（RV）模块搭建完成首条完整的语义分割流水线，代码位于Kaggle notebook，使用的数据集为NAIP-Chesapeake样本数据集。

训练后验证集预测结果表现正常，但使用以下代码为预测数据集生成预测并保存为GeoTIFF时，结果出现异常：

预测代码 #

from rastervision.core.data import SemanticSegmentationLabels

predictions = learner.predict_dataset(
    dataset=prediction_dataset,
    raw_out=True,
    numpy_out=True,
    predict_kw=dict(out_shape=IMG_SIZE), # 使用256x256
    progress_bar=True,
)

pred_labels = SemanticSegmentationLabels.from_predictions(
    windows=prediction_dataset.windows,
    predictions=predictions,
    smooth=True, # 尝试过开启/关闭该参数
    extent=prediction_dataset.scene.extent,
    num_classes=len(class_config),
)

保存GeoTIFF代码 #

pred_labels.save(
    uri=os.path.join(OUTPUT_DIR, 'predictions'),
    crs_transformer=prediction_dataset.scene.raster_source.crs_transformer,
    class_config=class_config,
)

结果对比 #

验证集预测结果：

预测保存前后对比（左侧为窗口显示的标签，右侧为保存到GeoTIFF文件中的同一窗口内容）：

可能的问题及解决方案 #

1. 预测输出格式与标签类预期不匹配 #

当设置raw_out=True时，模型输出的是未经过argmax处理的logits张量，但SemanticSegmentationLabels.from_predictions默认期望输入是类别ID（或正确处理的logits）。直接传入logits会导致保存时出现类别映射错误。

解决方法：

• 若保留raw_out=True，需手动对logits做argmax处理：

  # 对每个预测结果取argmax得到类别ID
  processed_predictions = [p.argmax(axis=0) for p in predictions]
  pred_labels = SemanticSegmentationLabels.from_predictions(
      windows=prediction_dataset.windows,
      predictions=processed_predictions,
      smooth=True,
      extent=prediction_dataset.scene.extent,
      num_classes=len(class_config),
  )
  

• 或者关闭raw_out=True，让predict_dataset直接返回类别ID：

  predictions = learner.predict_dataset(
      dataset=prediction_dataset,
      raw_out=False, # 修改为False
      numpy_out=True,
      predict_kw=dict(out_shape=IMG_SIZE),
      progress_bar=True,
  )
  

2. 窗口与预测结果尺寸不匹配 #

predict_kw=dict(out_shape=IMG_SIZE)设置的输出尺寸必须和prediction_dataset中窗口的实际尺寸严格对应，否则会导致拼接时出现错位、拉伸。

解决方法：

• 检查窗口尺寸是否与IMG_SIZE一致：

  print(prediction_dataset.windows[0].size)
  

• 若不一致，要么调整out_shape匹配窗口尺寸，要么重新生成固定为256x256的窗口。

3. 平滑参数的不合理使用 #

开启smooth=True时，会对重叠窗口的预测结果做加权平均，但如果输入是类别ID（而非logits），或者窗口重叠策略不合理，会导致输出模糊、类别混淆。

解决方法：

• 先尝试关闭smooth=False，查看保存结果是否恢复正常；
• 若需要平滑效果，确保输入的预测结果是logits，且窗口设置了合理的重叠率（如10%-20%）。

4. CRS坐标转换错误 #

保存GeoTIFF时使用的crs_transformer必须与预测场景的坐标系统完全一致，否则会出现空间错位，导致显示内容与窗口不匹配。

解决方法：

• 验证CRS是否匹配：

  print(prediction_dataset.scene.raster_source.crs_transformer.crs)
  

• 确保保存时使用的CRS与原始影像完全一致，避免坐标转换偏差。

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