构建图像训练数据列表时触发KeyError: '1'报错如何解决

我将图像数据集存放在 trainpath 文件夹中，trainpath目录下包含2个以类别名（0和1）命名的子文件夹，每个子文件夹下存储对应类别的图像数据，我尝试基于该数据集构建训练数据列表，编写的实现代码如下：

import os
import glob as gb
import cv2

# 注：s为提前定义好的图像缩放尺寸，trainpath为数据集根目录路径
code = {'distracted':0 ,'not-distracted':1}

X_train = []
y_train = []
for folder in  os.listdir(trainpath) : 
    files = gb.glob(pathname= str( trainpath +'//' + folder + '/*.jpg'))
    for file in files: 
        image = cv2.imread(file)
        image_array = cv2.resize(image , (s,s))
        X_train.append(list(image_array))
        y_train.append(code[folder])

运行上述代码时触发报错，报错信息如下：

KeyError                                  Traceback (most recent call last)
<ipython-input-74-5016fe834a0a> in <module>()
      7         image_array = cv2.resize(image , (s,s))
      8         X_train.append(list(image_array))
----> 9         y_train.append(code[folder])

KeyError: '1'

报错定位在y_train.append(code[folder])代码行。

触发KeyError的直接原因是字典键与实际遍历得到的文件夹名完全不匹配：

• 代码中定义的code字典键是'distracted'、'not-distracted'两个字符串
• 实际数据集下的子文件夹名是'0'、'1'，os.listdir()遍历到名为'1'的文件夹时，在code字典中找不到对应的键，直接抛出键错误。
• 额外注意：os.listdir()返回的所有文件名、文件夹名都是字符串类型，哪怕名称是纯数字，拿到的也是'0'、'1'格式的字符串，不是整数类型的0、1，写映射时要注意类型匹配。

直接修改code字典的键，和实际的子文件夹名保持一致即可，修正后的核心代码如下：

# 字典键改为和实际文件夹名完全一致的字符串'0'、'1'，值为需要对应生成的类别标签
code = {'0':0 ,'1':1}

X_train = []
y_train = []
for folder in  os.listdir(trainpath) : 
    files = gb.glob(pathname= str( trainpath +'//' + folder + '/*.jpg'))
    for file in files: 
        image = cv2.imread(file)
        # 增加判空逻辑，避免损坏图像读取失败触发后续resize报错
        if image is None:
            continue
        image_array = cv2.resize(image , (s,s))
        X_train.append(list(image_array))
        y_train.append(code[folder])

如果需要调整类别映射关系，比如把0文件夹对应为“非分心”、1文件夹对应为“分心”，只需要调换字典中键对应的标签值即可，核心要求是字典的键必须和遍历得到的文件夹名完全一致。

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