Python预处理函数传入不同目录参数却返回相同输出的问题

编写了一个preprocess函数，遍历CATEGORIES数组中命名的6个文件夹，将其中的图像转换为tensor后返回列表。但传入不同的目录参数调用该函数时，却得到了相同的输出。

相关代码 #

def preprocess(directory, img_size):
    X = []
    Y = []
    CATEGORIES = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']

    label = 0
    for category in CATEGORIES:
        path = os.path.join(DIRECTORY, category)
        for img in os.listdir(path):
            tensor = cv2.imread(os.path.join(path, img))
            tensor = cv2.resize(tensor, (img_size, img_size))
            X.append(tensor)
            Y.append(label)
        label +=1
    
    return X, Y

TRAIN_DIRECTORY = 'Data/seg_train'
TEST_DIRECTORY = 'Data/seg_test'

train_X, train_Y = preprocess(TRAIN_DIRECTORY, 150)
test_X, test_Y = preprocess(TEST_DIRECTORY, 150)

对比代码及输出 #

print("Train array size: ",len(train_Y), "\nTest array size: ",len(test_Y))

输出：

Train array size:  14034 
Test array size:  14034

问题原因 #

函数定义时接收的参数是小写的directory，但在拼接路径时错误使用了大写的DIRECTORY。这导致无论调用函数时传入什么目录参数，实际读取的都是DIRECTORY指向的同一个路径，两次调用自然会返回相同数量的样本。

修复方案 #

将路径拼接处的DIRECTORY替换为函数参数directory：

def preprocess(directory, img_size):
    X = []
    Y = []
    CATEGORIES = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']

    label = 0
    for category in CATEGORIES:
        # 改为使用函数传入的directory参数
        path = os.path.join(directory, category)
        for img in os.listdir(path):
            tensor = cv2.imread(os.path.join(path, img))
            tensor = cv2.resize(tensor, (img_size, img_size))
            X.append(tensor)
            Y.append(label)
        label +=1
    
    return X, Y

修改后，两次调用会分别读取TRAIN_DIRECTORY和TEST_DIRECTORY对应的路径，返回的样本数量会符合两个目录的实际文件数。

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