Kaggle之数据集的使用（2）类别存储在csv文件中

链接：Plant Seedlings Classification | Kaggle 

木薯叶病分类数据集中训练图片全部在一个文件夹中，并没有分类到各个对应的文件夹，采用train.csv来存储图片名称及对应的类别。

这时我们就不能将数据集划分，需要根据csv文件中的信息首先将每一类的图片存放到一个文件夹中。

import os
import pandas as pd
from torchvision.io import read_image
import PIL
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

class CustomImageDataset():
    #annotations_file指的是train_csv文本，img_dir指的是图片文件夹路径。
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        #img_labels是读取的csv，是一个dataframe文件，也就是表格
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, idx):
        #img_path指的是图片文件夹路径加上图片的名称，iloc[a,b],a是行索引，b是列索引
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = Image.open(img_path)
        #label是这张图片的真实标签
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label


def main():
    train_file = "train_images"
    img_dir = "train.csv"
    ci_leaf = CustomImageDataset(annotations_file=img_dir,img_dir=train_file)
    for i in range(len(ci_leaf)):
    #for i in range(2):
        image,label = ci_leaf.__getitem__(idx = i)
        print(label)
        s1 = str(i)
        s = s1 + ".jpg"
        if label == 0:
            #plt.show(image)
            image.save(os.path.join("./leav/CBB",s))
            #image.save("./leav/CBB/1.jpg")
        elif label == 1:
            image.save(os.path.join("./leav/CBSD", s))
            #image.save("./leav/CBSD")
        elif label == 2:
            image.save(os.path.join("./leav/CGM", s))
            #image.save("./leav/CGM")
        elif label == 3:
            image.save(os.path.join("./leav/CMD", s))
            #image.save("./leav/CMD/2.jpg")
        else :
            image.save(os.path.join("./leav/HEALTH", s))
            #image.save("./leav/HEALTH")


if __name__ == '__main__':
    main()

通过这个方法就将所有的图片分类到CBB,CBSD,CGM,CMD,HEALTH五个预先创建的文件夹中 。

再使用上一个文章使用的切分数据集的方法进行训练集和验证集的划分，就可以使用DataLoader方法读取数据集进行训练了。