Pytorch模型训练实用教程-代码解读(3)

最新推荐文章于 2023-05-04 21:36:16 发布

VIP文章 @@北冥有鱼

最新推荐文章于 2023-05-04 21:36:16 发布

阅读量407

点赞数

分类专栏： pytorch模型训练实用教程文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/m0_60096200/article/details/127430250

版权

# coding:utf-8
import os
'''
    为数据集生成对应的txt文件
'''

train_txt_path = os.path.join("..", "..", "Data", "train.txt")
train_dir = os.path.join("..", "..", "Data", "train")

valid_txt_path = os.path.join("..", "..", "Data", "valid.txt")
valid_dir = os.path.join("..", "..", "Data", "valid")


def gen_txt(txt_path, img_dir):
    f = open(txt_path, 'w')
    
    for root, s_dirs, _ in os.walk(img_dir, topdown=True):  # 获取 train文件下各文件夹名称
        for sub_dir in s_dirs:                               #  s_dirs是0-9
            i_dir = os.path.join(root, sub_dir)             # 获取各类的文件夹 绝对路径\train\0
            img_list = os.listdir(i_dir)                    # 获取类别文件夹下所有png图片的路径
            for i in range(len(img_list)):
                if not img_list[i].endswith('png