TensorFlow训练语义分割CNN时TIFF数据加载的可行方案问询

你遇到的报错核心原因是tf.data流水线传入的路径是张量类型，无法直接传入OpenCV、PIL这类接受字符串路径的Python库，而tfio的TIFF解码接口对自定义标签、特殊压缩格式的TIFF兼容性较差。目前最稳定的适配方案是使用TensorFlow官方提供的tf.numpy_function接口包装Python侧的TIFF加载逻辑，既可以兼容各类TIFF格式，也能完美适配tf.data生态。 #

具体实现代码 #

依赖导入 #

import tensorflow as tf
import cv2
import numpy as np

加载函数实现 #

def _load_tiff_data(img_path_bytes: bytes) -> tuple[np.ndarray, np.ndarray]:
    # 传入的路径是字节类型，转为字符串
    img_path = img_path_bytes.decode("utf-8")
    # 加载3通道RGB影像，cv2默认读取为BGR，需转格式
    image = cv2.imread(img_path, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 生成掩码路径，可根据自己的目录命名规则修改
    mask_path = img_path.replace("X", "y").replace("X.tif", "y.tif")
    # 加载单通道灰度掩码
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    # 扩展掩码维度为[H, W, 1]，和原生TF解码输出格式对齐
    mask = np.expand_dims(mask, axis=-1)
    return image.astype(np.uint8), mask.astype(np.uint8)

def parse_image(img_path: tf.Tensor) -> dict:
    # 用tf.numpy_function包装Python加载逻辑，适配tf.data流水线
    image, mask = tf.numpy_function(
        func=_load_tiff_data,
        inp=[img_path],
        Tout=[tf.uint8, tf.uint8]
    )
    # 显式设置张量形状，避免后续流水线形状推导报错
    # 如果你有固定的影像尺寸，可以直接写死例如[512, 512, 3]
    image.set_shape([None, None, 3])
    mask.set_shape([None, None, 1])
    return {"image": image, "segmentation_mask": mask}

使用注意事项 #

• 如果你的TIFF文件带有特殊压缩格式、地理坐标信息或者超过3通道，可以将cv2替换为PIL、rasterio等更专业的TIFF处理库，只需修改_load_tiff_data内部的加载逻辑即可，外层流水线不需要调整
• 性能优化：数据集较小的情况下可以在tf.data流水线中添加.cache()方法缓存数据，提升训练速度；大数据集可调整tf.data.Dataset.map的num_parallel_calls参数为tf.data.AUTOTUNE，开启多线程并行加载
• 如果你需要快速验证模型效果，也可以先批量将所有TIFF文件转成PNG格式，直接使用参考示例中的原生TF加载逻辑，无需修改适配代码

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