如何用Rasterio生态工具写入带5类Viridis色图的TIFF

关键修正与实现思路 #

• 数据离散化：先将原始数据映射为5个整数类别，确保像素值与色图颜色一一对应
• 色图格式适配：把Matplotlib的viridis色图转换为Rasterio支持的字典格式（键为类别值，值为RGBA四元组）
• 无显式GDAL依赖：完全依托rioxarray和rasterio的内置功能实现，无需直接导入GDAL

完整修改代码 #

import xarray as xr
import rioxarray
import os
import matplotlib.pyplot as plt
import pandas as pd

# 加载输入NC文件
ds_input = xr.open_dataset("inputfile.nc")
operatingdir = os.getcwd()  # 用当前工作目录避免硬编码路径问题
filename = os.path.splitext(os.path.basename("inputfile.nc"))[0]

bandcount = 7
# 创建5类Viridis色图并转为Rasterio兼容格式
cmap = plt.cm.get_cmap('viridis', 5)
colormap_dict = {
    i: tuple(int(c * 255) for c in cmap(i))
    for i in range(5)
}

for i in range(bandcount):
    # 获取当前波段数据（替换band_dim为实际的波段维度名，如time/step）
    band_data = ds_input['var_name'].isel(band_dim=i)
    
    # 将连续数据分箱为5个整数类别（0-4），保留nodata值
    bins = pd.cut(band_data.where(band_data != -9999), bins=5, labels=[0,1,2,3,4])
    binned_data = bins.astype('int16').fillna(-9999)
    
    # 输出带色图的云优化TIFF
    binned_data.rio.to_raster(
        os.path.join(operatingdir, "Output", f"{filename}_number{i}.tif"),
        windowed=True,
        compress='lzw',
        dtype='int16',
        nodata=-9999,
        overwrite=True,
        tiled=True,
        cloud_optimized=True,
        driver='COG',
        colormap=colormap_dict
    )

核心细节说明 #

1. 数据分箱：使用pd.cut将连续数值划分为5个区间，映射为0-4的整数类别。若你的数据本身已是离散类别值，可跳过此步骤，直接将类别值与色图字典的键对应。
2. 色图转换：遍历Matplotlib色图的5个颜色索引，将0-1范围的归一化颜色值转为0-255的整数RGBA元组，符合Rasterio的色图格式要求。
3. 波段索引：isel(band_dim=i)中的band_dim需替换为你NC文件中实际的波段维度名称（如time、forecast_step等），原代码未指定波段索引会导致输出多波段数据，需修正。

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