跟着小土堆学pytorch（p6.Dataset类代码实战）

from PIL import Image
img_path = "C:\\Users\\Administrator\\PycharmProjects\\learn_pytorch\\dataset\\train\\ants\\0013035.jpg"
img = Image.open(img_path)

img.size
Out[6]: (768, 512)

img.show()

dir_path = "dataset/train/ants"
import os
img_path_list = os.listdir(dir_path)

img_path_list[0]
Out[14]: '0013035.jpg'

f2 重命名

Dataset返回图片和标签，需要传入图片所在根目录和标签

1.PIL为自带的库 

2.img_path = "C:\\Users\\Administrator\\PycharmProjects\\learn_pytorch\\dataset\\train\\ants\\0013035.jpg"

为字符串变量

Image.open()传入图片地址，返回为PIL格式，如果用于torch训练需要totensor

3.绝对路径需要再加一个\转义，相对路径不用

4.img.show()可以查看图片

5.dir_path = "dataset/train/ants" 获取ants相对路径

6.import os   用os库中方法获取文件夹下目录树

7.img_path_list = os.listdir(dir_path)  获取ants下的图片名的数组

 8.获取第一张图片 img_path_list[0]
Out[14]: '0013035.jpg'

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset): #继承Dataset抽象类并重写init和len方法

    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)  # 调取path下的所有图片名称并返回一个数组


    def __getitem__(self,idx):
        img_name = self.img_path[idx]  # 通过init中定义的全局变量img_path 获取文件名
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name) # 将根目录+标签 + 文件名进行拼接形成相对地址
        img = Image.open(img_item_path)  #打开图片格式未PIL
        label = self.label_dir  # 标签就是self。lanel_dir 即 ants
        return img,label  # 返回 图片（PIL）格式，图片标签（label）

    def __len__(self):
        return len(self.img_path)

root_dir = "dataset/train"  # 定义root_dir 的 位置
ants_label_dir = "ants"     # 定义 ants_label_dir
ants_dataset = MyData(root_dir,ants_label_dir)  #  调用 自定义MyData类用于加载数据，返回图片和标签

9.root_dir 一般设置为

root_dir = "dataset/train" 根目录地址，字符串变量

10.label_dir = "ants" 这里的label就是ants

path = os.path.join(root_dir,label_dir)
path
Out[18]: 'dataset/train\\ants'

11.path = os.path.join(root_dir,label_dir) 将两个地址用\\拼接为一个新的字符串

root_dir = "dataset/train"  # 定义root_dir 的 位置
ants_label_dir = "ants"     # 定义 ants_label_dir
bees_label_dir = "bees" 
ants_dataset = MyData(root_dir,ants_label_dir)  #  调用 自定义MyData类用于加载数据，返回图片和标签
bees_dataset = MyData(root_dir,bees_label_dir)

 12.ants_dataset = MyData(root_dir,ants_label_dir)  初始化MyData，形成四个全局变量

13.ants_dataset[0] 返回第一个图片 和label故可以如下改写

img,label = ants_dataset[0]
img.show()

14.整个数据集：

train_dataset = ants_dataset + bees_dataset

两个数据集的集合

len(train_dataset)
Out[33]: 245
len(ants_dataset)
Out[34]: 124
len(bees_dataset)
Out[35]: 121 

img,label = train_dataset[123]
img.show()

 img,label = train_dataset[124]
img.show()

 适用于当数据集不足的时候仿造，添加数据集使用

批量写入标签：

import os

root_dir = 'dataset/train2'
target_dir = 'ants_image'
img_path = os.listdir(os.path.join(root_dir,target_dir))
label = target_dir.split('_')[0]  # label = ants 取分隔符前面部分
out_dir = 'ants_label'
for i in img_path:
    file_name = i.split('.jpg')[0] # file_name = .jpg 前面部分
    with open(os.path.join(root_dir,out_dir,"{}.txt".format(file_name)),'w') as f: # "{}.txt".format(file_name)) 文件名
        f.write(label)  # 写入标签ants

from torch.utils.data import Dataset
import cv2
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'
bees_label_dir = 'bees'
ants_dataset = MyData(root_dir,ants_label_dir)
bees_dataset = MyData(root_dir,bees_label_dir)
img,label = ants_dataset[8]
img.show()
img,label = bees_dataset[8]
img.show()