pytorch深度学习1：Dataset类

来自于 b 站《PyTorch 深度学习快速入门教程（绝对通俗易懂！）【小土堆】》

本篇内容主要是对代码中的一些小知识点进行归纳总结，以便复习

下面是主要代码

from torch.utils.data import Dataset    #引入Dataset类
from PIL import Image
import os

class MyData(Dataset):
    def __init__(self,root_dir,label_dir):  #属性

        self.root_dir=root_dir
        self.label_dir=label_dir
        self.path=os.path.join(self.root_dir,self.label_dir)
        self.img_path=os.listdir(self.path)     #返回路径指定的文件夹包含的文件或文件夹的名字的列表

    def __getitem__(self, idx):     #方法

        img_name=self.img_path[idx]
        img_item_path=os.path.join(self.root_dir,self.label_dir,img_name) #实现路径组合
        img=Image.open(img_item_path)
        label=self.label_dir
        return img, label

    def __len__(self):              #获取列表的长度
        return len(self.img_path)
#实际运行
root_dir="dataset/train"
ants_label_dir="ants_image"
ants_dataset=MyData(root_dir,ants_label_dir)
target_dir="ants_image"
label=target_dir.split('_')[0]
img_path=os.listdir(os.path.join(root_dir,target_dir))
out_dir="ants_label"
for i in img_path:
    file_name=i.split('.jpg')[0]
    with open(os.path.join(root_dir,out_dir,"{}.txt".format(file_name)),'w') as f:
        f.write(label)
        # 写文件操作，将标签写入root_dir/out_dir并创建txt文件，文件内容为label
        # 'w': 打开一个文件只用于写入。如果该文件已存在则将其覆盖。如果该文件不存在，创建新文件。

另外，配环境配了好久。。。