一、项目背景: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.textindex.upper()
return result
def len(self):
return len(self.text)
p = use_getitem('pytorch')
print(p1) # Y
print(len(p)) # 7
实现 getitem 后,实例能像列表一样用 p1 下标访问;实现 len 后,len(p) 能返回长度。为什么要先演示?因为自定义 Dataset 的机制就建立在这两个方法上 ------DataLoader 内部正是靠 len(dataset) 和 dataseti 取数据的。