使用os.listdir加载自定义图像数据集仅返回单个类别问题排查

问题说明 #

运行自定义create_dataset函数后，生成的target_dict仅包含{'forest': 0}，但测试时通过os.listdir能正常列出所有图像类别（forest、buildings、glacier、street、mountain、sea）。

错误原因 #

1. 返回语句缩进错误：create_dataset函数中的return img_data_array, class_name被放在了遍历类别目录的for dir1循环内部，导致函数处理完第一个有效类别（forest）后就直接返回，后续类别完全没被遍历。
2. 图像读取参数错误：cv2.imread的第二个参数传入cv2.COLOR_BGR2RGB是误用——该常量用于颜色空间转换，读取图像应使用cv2.IMREAD_COLOR，读取后再通过cv2.cvtColor转换为RGB格式。

修复后的完整代码 #

import pandas as pd
import numpy as np
import os
import tensorflow as tf
import cv2
from tensorflow import keras
import random
from tensorflow.keras.models import Sequential, Model
from  matplotlib import pyplot as plt
import matplotlib.image as mpimg
%matplotlib inline

IMG_WIDTH=200
IMG_HEIGHT=200
img_folder='//Users//julie-anne.rosquin//Documents//leongrosse//Intel_images//seg_train//seg_train'

def create_dataset(img_folder):
    img_data_array=[]
    class_name=[]
    for dir1 in os.listdir(img_folder):
        if dir1 != '.DS_Store':
            # 遍历当前类别下的所有图像文件
            for file in os.listdir(os.path.join(img_folder, dir1)):
                image_path= os.path.join(img_folder, dir1,  file)
                # 正确读取图像并转换为RGB格式
                image= cv2.imread(image_path, cv2.IMREAD_COLOR)
                image= cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                image=cv2.resize(image, (IMG_HEIGHT, IMG_WIDTH),interpolation = cv2.INTER_AREA)
                image=np.array(image)
                image = image.astype('float32')
                image /= 255 
                img_data_array.append(image)
                class_name.append(dir1)
    # 将return移到循环外部，确保遍历完所有类别后再返回
    return img_data_array, class_name

# 提取图像数组和类别名称
img_data, class_name =create_dataset(r'//Users//julie-anne.rosquin//Documents//leongrosse//Intel_images//seg_train//seg_train')

target_dict={k: v for v, k in enumerate(np.unique(class_name))}
print(target_dict)

预期输出 #

修复后运行代码，target_dict会包含所有6个类别：

{'buildings': 0, 'forest': 1, 'glacier': 2, 'mountain': 3, 'sea': 4, 'street': 5}

内容的提问来源于stack exchange，提问作者Julie-Anne