使用Pytorch框架自己制作做数据集进行图像分类（一）

第一章：Pytorch制作自己的数据集实现图像分类

第一章： Pytorch框架制作自己的数据集实现图像分类
第二章： Pytorch框架构建残差神经网络(ResNet)
第三章： Pytorch框架构建DenseNet神经网络

提示：本文代码，含有部分测试性输出语句，更改数据文件夹路径后可以直接跑通，文章末尾附全部代码

文章目录

• 第一章：Pytorch制作自己的数据集实现图像分类
• 前言
• 一、上网搜取相关照片作为数据
• 二、定义自己的数据类并读入图片数据
  *
• 1.引入相关库
• 2.继承Dataset实现Mydataset子类
• 3.使用glob方法获取文件夹中所有图片路径
• 三、为图片制作标签，并划分训练集与测试集
  *
• 1.利用自定义类Mydataset创建对象weather_dataset
• 2.为每张图片制作相应的标签
• 3.完善Mydataset类，将图片数据转换成Tensor，并展示部分图片与标签对应关系
• 4.划分数据集和测试集
• 总结
• 本文代码

前言

网上有很多直接利用已有数据集（如MNIST， CIFAR-10等），直接进行机器学习，图像分类的教程。但如何自己制作数据集，为图像制作相应标签等的教程较少。故写本文，分享一下自己利用Pytorch框架制作数据集的方法技巧。

开发环境：
Pycharm + Python 3.7.9
torch 1.10.2+cu102
torchvision 0.11.3+cu102

提示：以下是本篇文章正文内容

一、上网搜取相关照片作为数据

; 二、定义自己的数据类并读入图片数据

1.引入相关库

代码如下：

import glob
import torch
from torch.utils import data
from PIL import Image
import numpy as np
from torchvision import transforms
import matplotlib.pyplot as plt

2.继承Dataset实现Mydataset子类

代码如下：

class Mydataset(data.Dataset):

    def __init__(self, root):
        self.imgs_path = root

    def __getitem__(self, index):
        img_path = self.imgs_path[index]
        return img_path

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

init() 初始化方法，传入数据文件夹路径。
getitem() 切片方法，根据索引下标，获得相应的图片。
len() 计算长度方法，返回整个数据文件夹下所有文件的个数。

3.使用glob方法获取文件夹中所有图片路径

代码如下：

all_imgs_path = glob.glob(r'F:\weather\*\*.jpg')

for var in all_imgs_path:
    print(var)

三、为图片制作标签，并划分训练集与测试集

1.利用自定义类Mydataset创建对象weather_dataset

代码如下：

weather_dataset = Mydataset(all_imgs_path)
print(len(weather_dataset))
print(weather_dataset[12:15])
wheather_datalodaer = torch.utils.data.DataLoader(weather_dataset, batch_size=5)
print(next(iter(wheather_datalodaer)))

2.为每张图片制作相应的标签

代码如下：

species = ['cloud','sun','rain']
species_to_id = dict((c, i) for i, c in enumerate(species))
print(species_to_id)
id_to_species = dict((v, k) for k, v in species_to_id.items())
print(id_to_species)
all_labels = []

for img in all_imgs_path:

    for i, c in enumerate(species):
        if c in img:
            all_labels.append(i)
print(all_labels)

3.完善Mydataset类，将图片数据转换成Tensor，并展示部分图片与标签对应关系

代码如下：

transform = transforms.Compose([
                transforms.Resize((96,96)),
                transforms.ToTensor()
])

class Mydatasetpro(data.Dataset):

    def __init__(self, img_paths, labels, transform):
        self.imgs = img_paths
        self.labels = labels
        self.transforms = transform

    def __getitem__(self, index):
        img = self.imgs[index]
        label = self.labels[index]
        pil_img = Image.open(img)
        data = self.transforms(pil_img)
        return data, label

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

BATCH_SIZE = 10
weather_dataset = Mydatasetpro(all_imgs_path, all_labels, transform)
wheather_datalodaer = data.DataLoader(
                            weather_dataset,
                            batch_size=BATCH_SIZE,
                            shuffle=True
)

imgs_batch, labels_batch = next(iter(wheather_datalodaer))
print(imgs_batch.shape)

plt.figure(figsize=(12, 8))
for i, (img, label) in enumerate(zip(imgs_batch[:6], labels_batch[:6])):
    img = img.permute(1, 2, 0).numpy()
    plt.subplot(2, 3, i+1)
    plt.title(id_to_species.get(label.item()))
    plt.imshow(img)
plt.show()

4.划分数据集和测试集

代码如下：

index = np.random.permutation(len(all_imgs_path))

all_imgs_path = np.array(all_imgs_path)[index]
all_labels = np.array(all_labels)[index]

s = int(len(all_imgs_path)*0.8)
print(s)

train_imgs = all_imgs_path[:s]
train_labels = all_labels[:s]
test_imgs = all_imgs_path[s:]
test_labels = all_imgs_path[s:]

train_ds = Mydatasetpro(train_imgs, train_labels, transform)
test_ds = Mydatasetpro(test_imgs, test_labels, transform)
train_dl = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
test_dl = data.DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=True)

至此我们把原数据集的80%作为训练集得到了：
train_ds 训练集数据
test_ds 测试集数据
train_dl 训练集标签
test_dl 测试集标签

总结

整体思路
1.将自己所找图片按照一定的规则命名后，放到文件夹中。
2.使用glob方法获取所有数据文件路径
3.创建DataSet类的子类Mydataset，用于后续通过路径读入数据，并易于后续相应处理操作
4.通过Transforms.Compose()方法，对图片数据进行统一处理，并转换成Tensor格式
5.创建Mydatasetpro类，调用相关方法，获得{‘图片名’:标签}和{标签：’图片名’}字典
6.统计索引数量，按照百分比，划分出训练集和测试集
最后得到：训练集数据、测试集数据、训练集标签、测试集标签

本文代码

import glob
import torch
from torch.utils import data
from PIL import Image
import numpy as np
from torchvision import transforms
import matplotlib.pyplot as plt

class Mydataset(data.Dataset):

    def __init__(self, root):
        self.imgs_path = root

    def __getitem__(self, index):
        img_path = self.imgs_path[index]
        return img_path

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

all_imgs_path = glob.glob(r'F:\weather\*\*.jpg')

for var in all_imgs_path:
    print(var)

weather_dataset = Mydataset(all_imgs_path)
print(len(weather_dataset))
print(weather_dataset[12:14])
wheather_datalodaer = torch.utils.data.DataLoader(weather_dataset, batch_size=3)
print(next(iter(wheather_datalodaer)))

species = ['cloud','sun','rain']
species_to_id = dict((c, i) for i, c in enumerate(species))
print(species_to_id)
id_to_species = dict((v, k) for k, v in species_to_id.items())
print(id_to_species)
all_labels = []

for img in all_imgs_path:

    for i, c in enumerate(species):
        if c in img:
            all_labels.append(i)
print(all_labels)

transform = transforms.Compose([
                transforms.Resize((96,96)),
                transforms.ToTensor()
])

class Mydatasetpro(data.Dataset):

    def __init__(self, img_paths, labels, transform):
        self.imgs = img_paths
        self.labels = labels
        self.transforms = transform

    def __getitem__(self, index):
        img = self.imgs[index]
        label = self.labels[index]
        pil_img = Image.open(img)
        data = self.transforms(pil_img)
        return data, label

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

BATCH_SIZE = 10
weather_dataset = Mydatasetpro(all_imgs_path, all_labels, transform)
wheather_datalodaer = data.DataLoader(
                            weather_dataset,
                            batch_size=BATCH_SIZE,
                            shuffle=True
)

imgs_batch, labels_batch = next(iter(wheather_datalodaer))
print(imgs_batch.shape)

plt.figure(figsize=(12, 8))
for i, (img, label) in enumerate(zip(imgs_batch[:6], labels_batch[:6])):
    img = img.permute(1, 2, 0).numpy()
    plt.subplot(2, 3, i+1)
    plt.title(id_to_species.get(label.item()))
    plt.imshow(img)
plt.show()

index = np.random.permutation(len(all_imgs_path))

all_imgs_path = np.array(all_imgs_path)[index]
all_labels = np.array(all_labels)[index]

s = int(len(all_imgs_path)*0.8)
print(s)

train_imgs = all_imgs_path[:s]
train_labels = all_labels[:s]
test_imgs = all_imgs_path[s:]
test_labels = all_labels[s:]

train_ds = Mydatasetpro(train_imgs, train_labels, transform)
test_ds = Mydatasetpro(test_imgs, test_labels, transform)
train_dl = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
test_dl = data.DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)

Original: https://blog.csdn.net/zwy_697198/article/details/123561769
Author: 张_哈哈
Title: 使用Pytorch框架自己制作做数据集进行图像分类（一）