PyTorch自定义Dataset报TypeError：listdir路径类型错误的排查与解决

错误原因 #

1. __getitem__参数定义错误：PyTorch Dataset的__getitem__方法默认接收的是索引值（int类型），你错误地将参数命名为image_path，遍历dataset时传入的int索引被直接传给os.listdir()，触发了类型不匹配的错误。
2. __len__方法逻辑错误：你返回len(self.msr_data)，但self.msr_data从未在__init__中定义，属于未初始化的属性。
3. __getitem__核心逻辑偏离Dataset设计：当前写法每次调用__getitem__都会遍历整个文件夹的所有图片，而Dataset的核心是根据索引返回单样本，这种写法完全不符合预期。

修复方案 #

步骤1：在__init__中预加载所有图片路径 #

初始化阶段遍历目标文件夹，将所有图片的完整路径存入列表，避免每次获取样本时重复遍历文件夹，提升效率。

步骤2：修正__len__方法 #

返回预加载的图片路径列表的长度，让DataLoader能正确识别数据集的总样本数。

步骤3：修正__getitem__方法 #

• 将参数改为idx（索引值）
• 根据索引从路径列表中取出单张图片的路径
• 读取图片、转换为RGB格式、应用transform（如果有），最后返回单张图片数据

修正后的完整代码 #

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset

class msr_data(Dataset):
    
    def __init__(self, image_path, transform=None):
        self.image_path = image_path
        self.transform = transform
        # 预加载所有图片的完整路径，可选过滤非图片文件
        self.image_files = [
            os.path.join(image_path, filename) 
            for filename in os.listdir(image_path)
            if filename.lower().endswith(('.png', '.jpg', '.jpeg'))
        ]

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        img_path = self.image_files[idx]
        img = cv2.imread(img_path)
        im_rgb = img[:, :, ::-1]
        if self.transform is not None:
            im_rgb = self.transform(im_rgb)
        return np.array(im_rgb)
 
dataset = msr_data(image_path='files', transform=None)
for i, sample in enumerate(dataset):
    print(f"样本索引: {i}")
    plt.imshow(sample)
    plt.axis('off')
    plt.show()

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