如何在TensorFlow数据管道中用cv2将tif图像路径转换为数组？

报错核心原因 #

• tf.py_function 传入to_array的路径参数是tf.Tensor类型，未转换为Python字符串，cv2.imread无法识别张量作为输入路径，返回NULL导致后续操作报错
• py_function的输出张量未手动指定形状，TensorFlow无法自动推导张量维度，导致输出shape显示为
• 未对路径有效性做校验，存在错误路径或损坏图片时会无明确提示直接崩溃

修复后的完整代码 #

import cv2
import numpy as np
import tensorflow as tf

# 全局定义图像尺寸，可根据需求修改
IMG_H = 256
IMG_W = 256

def to_array(p_0):
    # 将tf张量转为Python字符串路径
    path = p_0.numpy().decode('utf-8')
    im_1 = cv2.imread(path, cv2.IMREAD_COLOR)
    # 增加路径有效性校验
    if im_1 is None:
        raise ValueError(f"无法读取图像路径：{path}，请检查路径是否存在或图像是否损坏")
    im_1 = cv2.resize(im_1, (IMG_W, IMG_H))
    im_1 = np.asarray(im_1, dtype=np.float64)
    im_1 /= 255.0
    return im_1

def parse_fn(p_0):
    [im,] = tf.py_function(to_array, [p_0], [tf.float64])
    # 手动设置张量形状，让TensorFlow可以推导维度
    im.set_shape((IMG_H, IMG_W, 3))
    return im

def TFData(p_0, p_1, batch_size, labels=None, is_train=True):
    dset_1 = tf.data.Dataset.from_tensor_slices(p_0)
    dset_1 = dset_1.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
    dset_2 = tf.data.Dataset.from_tensor_slices(p_1)
    dset_2 = dset_2.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)

    if labels is None:
        raise ValueError("labels参数不能为空")
    label_dset = tf.data.Dataset.from_tensor_slices(labels)
  
    # 按照需求打包为((img1, img2), label)的结构
    final_dset = tf.data.Dataset.zip(((dset_1, dset_2), label_dset))
    final_dset = final_dset.batch(batch_size, drop_remainder=is_train).prefetch(tf.data.AUTOTUNE)
    return final_dset

性能优化建议 #

如果想要规避tf.py_function带来的GIL锁性能损耗，可以用TensorFlow原生API实现图像读取，流水线运行效率更高：

import tensorflow_io as tfio

def parse_fn_tf(path):
    img = tf.io.read_file(path)
    # 解码tif图像
    img = tfio.experimental.image.decode_tiff(img)
    # 裁剪掉tif可能带的alpha通道，保留3通道
    img = img[..., :3]
    img = tf.image.resize(img, (IMG_H, IMG_W))
    img = tf.cast(img, tf.float64) / 255.0
    return img

替换原有parse_fn即可，不需要再用py_function包裹。

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