在TensorFlow的map方法中使用Rawpy读取RAW图像报错求助

为啥会报错 #

• 直接在map里用rawpy.imread()：tf.data的map默认以图模式运行，传入的是Tensor对象，但rawpy.imread()只接受普通字符串路径，Tensor无法被直接解析，因此触发'Tensor' object has no attribute 'encode'错误。
• 尝试用.numpy()提取字符串：图模式下的Tensor不支持直接调用.numpy()，只有在eager执行模式或通过tf.py_function包装的函数中才能访问Tensor的numpy值，所以又出现了第二个错误。

怎么解决 #

使用tf.py_function将读取RAW图像的函数包装起来，让它能在图模式中处理Tensor，并返回Tensor格式的数据。修改后的代码如下：

import tensorflow as tf
import rawpy
import numpy as np

def read_image(im1_path, im2_path):
    # 将Tensor转为Python字符串
    im1_path_str = im1_path.numpy().decode('utf-8')
    im2_path_str = im2_path.numpy().decode('utf-8')
    
    # 读取RAW图像并转换为float32格式
    with rawpy.imread(im1_path_str) as raw1:
        im1 = raw1.raw_image_visible.astype(np.float32)
    with rawpy.imread(im2_path_str) as raw2:
        im2 = raw2.raw_image_visible.astype(np.float32)
    
    # 将numpy数组转回Tensor格式
    return tf.convert_to_tensor(im1), tf.convert_to_tensor(im2)

# 用tf.py_function包装函数，指定输入输出的数据类型
ds_train = ds_train.map(lambda x, y: tf.py_function(
    func=read_image,
    inp=[x, y],
    Tout=[tf.float32, tf.float32]
))

关键细节 #

• tf.py_function的作用是在图模式中执行Python函数，它会自动把输入的Tensor转为numpy数组，执行完函数后再转回Tensor。
• 必须指定Tout参数，明确函数返回值的数据类型，确保后续训练流程能正确识别数据集的类型。
• 使用with rawpy.imread(...)的写法可以自动释放资源，避免内存泄漏问题。

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