运行ViT模型遇Tensor转换错误：PIL图像转Tensor失败求解

问题 #

运行《Image Classification with Hugging Face Transformers and Keras》中的ViT模型时，执行代码processed_dataset = ds.map(augmentation, batched=True)出现以下错误：

ValueError: Exception encountered when calling layer "resizing_8" (type Resizing).
Attempt to convert a value (<PIL.BmpImagePlugin.BmpImageFile image mode=L size=190x100 at 0x7F35C52AD210>) with an unsupported type (<class 'PIL.BmpImagePlugin.BmpImageFile'>) to a Tensor.
Call arguments received: • inputs=<PIL.BmpImagePlugin.BmpImageFile image mode=L size=190x100 at 0x7F35C52AD210>

尝试在create_image_folder_dataset()的features中添加'img': Image(decode=True, id=None)后，仅layer名称变为resizing_13，错误仍未解决。

相关代码 #

数据集创建函数

def create_image_folder_dataset(root_path):
  """creates `Dataset` from image folder structure"""
  
  # get class names by folders names
  _CLASS_NAMES= os.listdir(root_path)
  # defines `datasets` features`
  features=datasets.Features({
                      "img": datasets.Image(decode=True, id=None),
                      #"img": datasets.Image(),
                      "label": datasets.features.ClassLabel(names=_CLASS_NAMES),
                       
                  })
  #print(_CLASS_NAMES)
  # temp list holding datapoints for creation
  img_data_files=[]
  label_data_files=[]
  # load images into list for creation
  for img_class in os.listdir(root_path):
    for img in os.listdir(os.path.join(root_path,img_class)):
      path_=os.path.join(root_path,img_class,img)
      img_data_files.append(path_)
      label_data_files.append(img_class)
  # create dataset
  ds = datasets.Dataset.from_dict({"img":img_data_files,"label":label_data_files},features=features)
  return ds
ds = create_image_folder_dataset("/content/drive/MyDrive/FINAL_DATASET")
ds[0] """ return: 
{'img': <PIL.BmpImagePlugin.BmpImageFile image mode=L size=190x100 at 0x7F35C54ECC10>,
 'label': 0}"""

数据增强函数

from transformers import ViTFeatureExtractor
from tensorflow import keras 
from tensorflow.keras import layers


model_id = "google/vit-base-patch16-224-in21k"
#google/vit-base-patch32-384
feature_extractor = ViTFeatureExtractor.from_pretrained(model_id)

# learn more about data augmentation here: https://www.tensorflow.org/tutorials/images/data_augmentation
data_augmentation = keras.Sequential(
    [
        layers.Resizing(feature_extractor.size, feature_extractor.size),
        layers.Rescaling(1./255),
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(factor=0.02),
        layers.RandomZoom(
            height_factor=0.2, width_factor=0.2
        ),
    ],
    name="data_augmentation",
)
# use keras image data augementation processing
def augmentation(examples):
    print(examples["img"])
    examples["pixel_values"] = [data_augmentation(image) for image in examples["img"]]
    return examples


# basic processing (only resizing)
def process(examples):
    examples.update(feature_extractor(examples['img'], ))
    return examples


# we are also renaming our label col to labels to use `.to_tf_dataset` later
#ds = ds.rename_column("label", "labels")

解决方案 #

错误原因 #

Keras的预处理层（如Resizing）无法直接处理PIL图像对象，必须先转换为numpy数组或TensorFlow张量；同时你的图像是单通道灰度图（mode=L），而ViT预训练模型默认要求3通道RGB输入，这也会引发后续适配问题。

具体修复 #

修改augmentation函数，先将PIL图像转为numpy数组，处理单通道转3通道的问题，同时改用批量处理提升效率：

import numpy as np

def augmentation(examples):
    processed_images = []
    for img in examples["img"]:
        # 将PIL图像转为numpy数组
        img_np = np.array(img)
        # 单通道灰度图扩展为3通道，匹配ViT输入要求
        if len(img_np.shape) == 2:
            img_np = np.repeat(img_np[..., np.newaxis], 3, axis=-1)
        processed_images.append(img_np)
    
    # 批量传入数据增强流水线，比单张处理效率更高
    examples["pixel_values"] = data_augmentation(np.array(processed_images))
    return examples

额外优化 #

如果不需要保留原img字段，可以在map时指定remove_columns=["img"]，减少内存占用：

processed_dataset = ds.map(augmentation, batched=True, remove_columns=["img"])

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