基于YouTube教程搭建AI图像分类器时遇NotADirectoryError问题求助

一、修复NotADirectoryError错误 #

错误原因 #

你的train_dir直接指向了存放图片的Blue-Squares文件夹，val_dir指向Red-Triangles文件夹。当执行for category in os.listdir(dir_)时，category获取到的是文件夹里的图片文件名，而非子目录名。后续代码试图把图片文件当作目录去遍历，因此抛出“目录名称无效”的错误。

解决方案 #

调整代码逻辑，直接遍历训练/验证目录下的图片文件，无需嵌套子目录遍历：

from img2vec_pytorch import Img2Vec
import os
from PIL import Image

# prepare data

# 解决torchvision警告：指定weights参数
img2vec = Img2Vec(model='resnet18', weights='DEFAULT')

data_dir = 'C:/XX/XX/XX/XX/newai'
train_dir = os.path.join(data_dir, 'Data', 'Blue-Squares')
val_dir = os.path.join(data_dir, 'Data', 'Red-Triangles')

data = {}

for j, dir_ in enumerate([train_dir, val_dir]):
    features = []
    labels = []
    # 直接遍历当前目录下的所有文件
    for img_name in os.listdir(dir_):
        img_path_ = os.path.join(dir_, img_name)
        # 过滤非图片文件（避免处理无关文件）
        if not img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
            continue
        try:
            img = Image.open(img_path_)
            img_features = img2vec.get_vec(img)
            features.append(img_features)
            # 用当前文件夹名作为标签（比如Blue-Squares/Red-Triangles）
            labels.append(os.path.basename(dir_))
        except Exception as e:
            print(f"处理图片{img_path_}失败: {e}")
            continue

    data[['training_data', 'validation_data'][j]] = features
    data[['training_labels', 'validation_labels'][j]] = labels


print(data.keys())

# train model

# test performance

# save the model

关键修改点：

• 移除嵌套的子目录遍历循环，直接读取目标目录下的图片
• 用文件夹名作为对应数据集的统一标签
• 添加图片格式过滤和异常捕获，提升代码鲁棒性

二、消除torchvision弃用警告 #

警告原因 #

img2vec_pytorch内部调用的torchvision模型仍在使用已弃用的pretrained参数，torchvision 0.13+版本推荐使用weights参数替代。

解决方案 #

初始化Img2Vec时显式指定weights参数：

img2vec = Img2Vec(model='resnet18', weights='DEFAULT')

'DEFAULT'会自动加载对应模型的最新权重，也可以指定具体权重枚举（如ResNet18_Weights.IMAGENET1K_V1），但'DEFAULT'更简洁通用。

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