解决os/shutil的FileNotFoundError及PyTorch数据集分类问题

问题背景 #

我有如下数据集结构：

• train/test 文件夹，各自包含：
  • images 子文件夹：存储JPEG图片（如image_0000.jpg）
  • labels 子文件夹：存储对应TXT标签文件（如image_0000.txt，文件内容为0/1/2中的一个类别）
    已创建train_pytorch/test_pytorch文件夹，并在其中建立0、1、2类别子文件夹，需要将图片按对应标签移动到对应类别子文件夹以适配PyTorch的ImageFolder。

原代码运行时出现FileNotFoundError，提示路径C:\Users\rgs88\OneDrive\Documents\Semester 2 Data Science Masters\Computer Vision\CV2023_CW_Dataset\train/images不存在，且疑惑路径中为何显示正斜杠。

问题分析 #

1. 路径显示问题：报错中的正斜杠是Python打印路径时的统一表现，os.path.join在Windows下会自动生成正确的反斜杠路径，报错核心是目标路径确实不存在，需先确认train/images、train/labels等子文件夹是否真实存在，路径拼写是否正确。
2. 逻辑错误：原代码误将标签文件名后缀当作类别判断依据，但实际类别存储在TXT文件内容中，而非文件名。
3. 循环冗余：原代码每次遍历类别时都重复读取所有图片和标签文件，效率低下。

解决方案与修正代码 #

修正思路 #

• 新增路径存在性检查，提前发现路径错误
• 通过读取TXT文件内容获取真实类别，而非依赖文件名
• 用函数封装逻辑，同时处理train和test数据集，减少冗余
• 加入异常处理：跳过不存在的图片或无效标签

完整修正代码 #

import os
import shutil

# 定义所有路径
train_dir = r"C:\Users\rgs88\OneDrive\Documents\Semester 2 Data Science Masters\Computer Vision\CV2023_CW_Dataset\train"
test_dir = r"C:\Users\rgs88\OneDrive\Documents\Semester 2 Data Science Masters\Computer Vision\CV2023_CW_Dataset\test"
pytorch_train_dir = r"C:\Users\rgs88\OneDrive\Documents\Semester 2 Data Science Masters\Computer Vision\CV2023_CW_Dataset\train_pytorch"
pytorch_test_dir = r"C:\Users\rgs88\OneDrive\Documents\Semester 2 Data Science Masters\Computer Vision\CV2023_CW_Dataset\test_pytorch"
class_names = ["0", "1", "2"]

def organize_dataset(source_dir, target_dir):
    # 检查源路径是否存在
    images_dir = os.path.join(source_dir, "images")
    labels_dir = os.path.join(source_dir, "labels")
    assert os.path.exists(images_dir), f"路径不存在：{images_dir}"
    assert os.path.exists(labels_dir), f"路径不存在：{labels_dir}"

    # 提前创建所有类别文件夹
    for cls in class_names:
        os.makedirs(os.path.join(target_dir, cls), exist_ok=True)

    # 遍历所有标签文件
    for label_filename in os.listdir(labels_dir):
        if not label_filename.endswith(".txt"):
            continue
        
        # 获取对应图片文件名
        img_filename = label_filename.replace(".txt", ".jpg")
        img_path = os.path.join(images_dir, img_filename)
        
        # 检查图片是否存在
        if not os.path.exists(img_path):
            print(f"警告：图片 {img_filename} 不存在，跳过")
            continue
        
        # 读取标签内容（获取类别）
        label_path = os.path.join(labels_dir, label_filename)
        with open(label_path, "r") as f:
            cls = f.read().strip()
            # 确保类别是预期的0/1/2
            if cls not in class_names:
                print(f"警告：标签 {label_filename} 内容无效：{cls}，跳过")
                continue
        
        # 复制图片到对应类别文件夹（如需移动文件，用shutil.move替代shutil.copyfile）
        target_img_path = os.path.join(target_dir, cls, img_filename)
        shutil.copyfile(img_path, target_img_path)
        # 可选：复制标签文件到对应文件夹（ImageFolder不需要标签文件，若需保留可取消注释）
        # target_label_path = os.path.join(target_dir, cls, label_filename)
        # shutil.copyfile(label_path, target_label_path)

# 处理训练集
organize_dataset(train_dir, pytorch_train_dir)
# 处理测试集
organize_dataset(test_dir, pytorch_test_dir)

print("数据集整理完成")

代码关键点说明 #

• 若需要移动文件而非复制，将shutil.copyfile替换为shutil.move即可
• PyTorch的ImageFolder仅需要图片文件，标签文件可根据需求选择是否保留
• 代码中加入了路径检查和异常处理，避免因文件缺失或无效标签导致程序崩溃

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