一、项目背景:20 类食物,怎么"喂"给模型?

目标是训练一个卷积神经网络(CNN),识别 20 类食物:八宝粥、哈密瓜、圣女果、巴旦木、板栗、汉堡、火龙果、炸鸡、瓜子、生肉、白萝卜、胡萝卜、草莓、菠萝、薯条、蛋、蛋挞、青菜、骨肉相连、鸡翅。

数据按"一个文件夹一个类别"的方式组织,这是图像分类数据集的标准结构:

数据集目录结构(food_dataset)

food_dataset/

├── train/                  # 训练集

│   ├── 八宝粥/             # 文件夹名 = 类别名

│   ├── 哈密瓜/

│   └── ...                 # 共 20 个类别文件夹

└── test/                   # 测试集,结构相同

关键约定:文件夹名就是类别名;数据已按 train(训练集)/test(测试集)分好。写网络之前的第一步,是把"图片路径 + 类别编号"整理成清单文件,让 PyTorch 按行读取——这就是数据处理.py 的任务。

二、核心函数:用 os.walk 遍历目录树

数据处理.py —— 生成标签清单的核心函数

def train_test_file(root,dir):

    file_txt = open(dir+'.txt','w',encoding='utf-8')

    path = os.path.join(root,dir)

    for roots,directories,files in os.walk(path):

        if len(directories) !=0 :

            dirs = directories

        else:

            now_dir = roots.split('\\')

            for file in files:

                path_1 = os.path.join(roots,file)

                file_txt.write(path_1+' '+str(dirs.index(now_dir[-1]))+'\n')

    file_txt.close()

函数作用:遍历给定目录,把每张图片写进 txt,格式为「图片路径 + 空格 + 类别编号」,一行一条:

.\food_dataset\train\八宝粥\img_八宝粥罐_22.jpeg 0

.\food_dataset\train\哈密瓜\img_哈密瓜_103.jpeg 1

• os.path.join(root, dir) 拼出要扫描的目录(如 ./food_dataset/train);

• os.walk 逐层遍历,每次返回三元组:roots(当前目录)、directories(子文件夹)、files(文件);

• walk "先父后子":第一次进入 train 时 directories 就是全部 20 个类别文件夹,先存进 dirs 备用;

• 走到叶子目录(不再有子文件夹)时,把里面每张图的路径和类别编号写入文件。

三、重点代码:文件夹名如何变成数字标签

标签映射 —— 全函数最关键的两行

now_dir = roots.split('\\')            # 按 \ 切开路径,取最后一段

file_txt.write(path_1+' '+str(dirs.index(now_dir[-1]))+'\n')

roots.split('\\') 把路径按反斜杠切开取最后一段——类别文件夹名(如"八宝粥");dirs.index(该名字) 返回它在 20 个类别列表中的下标 0~19,这个下标就是数字标签。训练/测试两次调用文件夹顺序一致,标签能对齐;str(...) 把数字转成字符串,'\n' 换行。

四、必须强调的坑:中文乱码

open(dir+'.txt','w',encoding='utf-8') 里的 encoding='utf-8' 不能省。Windows 中文系统默认用 GBK 编码写文件,而 PyTorch 读取 txt 按 UTF-8 解码——类别名是中文,读出来全是乱码,后续 Image.open 找不到文件直接报错。写、读两端统一 UTF-8 最省事(第 2 篇读文件同样要指定)。

五、铺垫:两个魔法方法

数据处理.py —— 魔法方法演示

class use_getitem():

    def __init__(self,text):

        self.text = text

    def __getitem__(self,index):

        result = self.text[index].upper()

        return result

    def __len__(self):

        return len(self.text)

 

p = use_getitem('pytorch')

print(p[1])      # Y

print(len(p))    # 7

实现 __getitem__ 后,实例能像列表一样用 p[1] 下标访问;实现 __len__ 后,len(p) 能返回长度。为什么要先演示?因为自定义 Dataset 的机制就建立在这两个方法上——DataLoader 内部正是靠 len(dataset) 和 dataset[i] 取数据的。