如何使用自定义map()函数从GeoTiff文件创建TF Dataset

tf.data.Dataset的map算子默认运行在TensorFlow图模式下，传入parse_image的img_path参数并非原生Python字符串，而是tf.string类型的张量，rasterio的open方法无法识别TensorFlow张量类型的输入，仅支持原生字符串、路径对象或文件类对象，因此触发类型错误。

使用tf.py_function将调用rasterio的Python原生逻辑包装为TensorFlow图可识别的操作，该算子会自动将输入张量转换为numpy数组，你只需要在函数内将路径张量转换为原生Python字符串即可。

修改后的可运行代码如下：

import tensorflow as tf
import rasterio as rio

def parse_image(img_path: str) -> dict:
    # 把Tensor转换为原生Python字符串
    img_path = img_path.numpy().decode('utf-8')
    with rio.open(img_path) as src:
        image_numpy = src.read()
    image = tf.convert_to_tensor(image_numpy, dtype=tf.float32)
    return {'image': image}

def wrap_parse_function(img_path):
    # 用tf.py_function包装原生Python函数，指定输入输出类型
    output = tf.py_function(func=parse_image, inp=[img_path], Tout={'image': tf.float32})
    # 可选：手动设置张量shape，避免后续计算图推断shape失败
    # 假设你的tiff是[C, H, W]格式，根据实际维度修改即可
    # output['image'].set_shape([num_channels, height, width])
    return output

train_dataset = tf.data.Dataset.list_files("multipage_tiff_example.tif", seed=50)
train_dataset = train_dataset.map(wrap_parse_function)

• 如果需要更高的加载性能，可以在数据集构造链路后添加prefetch(tf.data.AUTOTUNE)，或搭配interleave算子实现并行加载
• 若你的tif文件尺寸固定，建议手动调用set_shape设置张量的维度信息，避免后续模型搭建时出现维度未知的报错
• tf.py_function的逻辑只能在CPU上运行，若要进一步优化加载效率，可以提前将所有tiff文件读取后序列化存储为TFRecord格式，训练时直接读取TFRecord即可避免训练过程中反复调用rasterio的开销

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