自制数据集

一 数据集

下载链接： http://download.tensorflow.org/example_images/flower_photos.tgz

打开后里面有5个文件夹，下图是我的文件目录：

下图是pycharm目录

二 整体思路

1.以字典的形式给5个文件夹制定标签，daisy:0 dandelion:1 等

2.分别读取5个文件夹中的图片名，使得图片名和文件夹便签相对应，并保存在images.csv文件中

如：‘I:/flowers/flower_photos\tulips\8908097235_c3e746d36e_n.jpg’， 4

3.读取images.csv制作数据和标签文件

4.处理数据和标签文件

三 generate_data.py

1.先把要用到的包加载一下

import os
import csv
import glob
import torch
import numpy as np
import pandas as pd
from PIL import Image
from torchvision import transforms

2.以字典的形式给5个文件夹制定标签

name2label = {}
path = "I:\\flowers_data\\flower_photos"
for name in (os.listdir(path)):
    # 为空则运行continue
    if not os.path.isdir(os.path.join(path, name)):
        continue
    # 以字典形式制定标签
    name2label[name] = len(name2label.keys())
#  {'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}   

3.读取图片名并保存在images.csv中

# 读取文件名
images = []
for name in name2label.keys():
    images += glob.glob(os.path.join(path, name, '*.png'))
    images += glob.glob(os.path.join(path, name, '*.jpg'))
    images += glob.glob(os.path.join(path, name, '*.jpeg'))
# 'I:/flowers/flower_photos\\tulips\\8908097235_c3e746d36e_n.jpg'

# 保存在images.csv中
    with open('images.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        for img in images:
            # 得到I:/flowers/flower_photos\\tulips\\8908097235_c3e746d36e_n.jpg里的倒数第二项，即花名tulips
            name = img.split(os.sep)[-2]
            # 根据花名得到标签名
            label = name2label[name]
            writer.writerow([img, label])
# I:\flowers_data\flower_photos\tulips\5700394524_dc6f8fa9cd_n.jpg,4            

4.制作图片和标签文件

images, labels = [], []
with open(r"images.csv") as f:
    reader = csv.reader(f)
    for row in reader:
        img, label = row
        label = int(label)
        images.append(img)
        labels.append(label)

5.处理图片和标签文件

image_data = []
for i in range(len(images)):
    tf = transforms.Compose([
                lambda x:Image.open(x).convert('RGB'),   
                transforms.Resize((224, 224)),
                transforms.ToTensor()
        ])
    img = tf(images[i])
    image_data.append(img)
    
# image_data是多维数据，没法像labels一样直接转换
all_X = torch.tensor([item.cpu().detach().numpy() for item in image_data])
# print(type(all_X))  <class 'torch.Tensor'>
all_Y = np.array(labels)
# print(type(all_Y))   <class 'numpy.ndarray'>

np.save("all_X", all_X)
np.save("all_Y", all_Y)

6.全部代码

import os
import csv
import glob
import torch
import numpy as np
import pandas as pd
from PIL import Image
from torchvision import transforms

name2label = {}
path = "I:\\flowers_data\\flower_photos"
for name in (os.listdir(path)):
    # 为空则运行continue
    if not os.path.isdir(os.path.join(path, name)):
        continue
    name2label[name] = len(name2label.keys())
# print(name2label)
# 读取文件名
images = []
for name in name2label.keys():
    images += glob.glob(os.path.join(path, name, '*.png'))
    images += glob.glob(os.path.join(path, name, '*.jpg'))
    images += glob.glob(os.path.join(path, name, '*.jpeg'))
# 'I:/flowers/flower_photos\\tulips\\8908097235_c3e746d36e_n.jpg'

# 保存成csv文件
    with open('images.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        for img in images:
            # 得到文件名里的倒数第二项，即花名
            name = img.split(os.sep)[-2]
            label = name2label[name]
            writer.writerow([img, label])

images, labels = [], []
with open(r"images.csv") as f:
    reader = csv.reader(f)
    for row in reader:
        img, label = row
        label = int(label)
        images.append(img)
        labels.append(label)

# print(len(images), len(labels))
image_data = []
for i in range(len(images)):
    tf = transforms.Compose([
                lambda x:Image.open(x).convert('RGB'),   
                transforms.Resize((224, 224)),
                transforms.ToTensor()
        ])
    # img = tf('I:\\flowers\\flower_photos\daisy\\5547758_eea9edfd54_n.jpg')
    img = tf(images[i])
    image_data.append(img)


# image_data是多维数据，没法像labels一样直接转换
all_X = torch.tensor([item.cpu().detach().numpy() for item in image_data])
# print(type(all_X))  <class 'torch.Tensor'>
all_Y = np.array(labels)
# print(type(all_Y))   <class 'numpy.ndarray'>

np.save("all_X", all_X)
np.save("all_Y", all_Y)

四 main.py

import torch
import time
from torch import nn
from torch.utils.data import Dataset, DataLoader, TensorDataset
import numpy as np
from sklearn.model_selection import train_test_split
import torch.optim as optim

batch_size = 16
learning_rate = 0.001
weight_decay = 6e-5
episode_num = 100

# 数据加载
X = np.load("all_X.npy", allow_pickle=True)
# x,shape:(3670, 3, 224, 224)
Y = np.load("all_Y.npy", allow_pickle=True)
# (3670,)

# 数据集划分
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.20, random_state=99)
# X_train.shape:(2936, 3, 224, 224)
# y_train.shape:(2936,)

# 把数组转换成张量，且二者共享内存
X_train = torch.from_numpy(X_train)
y_train = torch.from_numpy(y_train)
X_test = torch.from_numpy(X_test)
y_test = torch.from_numpy(y_test)

# 对给定的tensor数据(样本和标签)，将它们包装成dataset。如果是numpy的array，或者Pandas的DataFrame需要先转换成Tensor
data = TensorDataset(X_train, y_train)
train_loader = DataLoader(data, batch_size=batch_size, shuffle=False)

data = TensorDataset(X_test, y_test)
test_loader = DataLoader(data, batch_size=batch_size, shuffle=False)


# 模型定义
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, 3, 1)
        self.relu1 = nn.ReLU()
        self.bn1 = nn.BatchNorm2d(8)
        self.dropout1 = nn.Dropout()
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(8, 16, 3, 1)
        self.relu2 = nn.ReLU()
        self.bn2 = nn.BatchNorm2d(16)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.dropout2 = nn.Dropout()
        self.out = nn.Linear(16*54*54, 5)

    def forward(self, x):
        # print(x.shape)
        # torch.Size([16, 3, 224, 224])
        out = self.bn1(self.relu1(self.conv1(x)))
        out = self.pool1(self.dropout1(out))
        out = self.bn2(self.relu2(self.conv2(out)))
        out = self.pool2(self.dropout2(out))
        out = out.view(out.size(0), -1)
        out = self.out(out)
        return out


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model().to(device)

optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
criterion = nn.CrossEntropyLoss()
criterion.to(device)

# 释放显存
torch.cuda.empty_cache()
for episode in range(episode_num):
    start = time.perf_counter()
    print("===Episode {}/{}===".format(episode+1, episode_num))
    train_acc = 0
    train_loss = 0
    model.train()
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        out = model(batch_x.float())
        loss = criterion(out, batch_y.long())
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 记录误差
        train_loss += loss.item()

        # 计算分类的准确率
        _, pred_y = out.max(1)
        num_correct = (pred_y == batch_y).sum().item()
        acc = num_correct / batch_x.size(0)
        train_acc += acc

    # 在测试集上检验效果
    eval_loss = 0
    eval_acc = 0

    # 在评估模型时使用，固定BN 和 Dropout
    model.eval()
    for img, label in test_loader:
        img, label = img.to(device), label.to(device)
        # img = img.view(img.size(0), -1)
        # 前向传播
        out = model(img.float())
        loss = criterion(out, label.long())
        eval_loss += loss.item()
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        eval_acc += acc
    end = time.perf_counter()
    print('Train Loss: {:.6f}, Train Acc: {:.6f}, Test Loss:{:.6f}, Acc:{:.6f}, Running time:{:.4f} s'
          .format(train_loss / len(train_loader), train_acc / len(train_loader), eval_loss / len(test_loader),
                  eval_acc / len(test_loader), (end-start)))

参考博客

https://blog.csdn.net/hxxjxw/article/details/106301977?utm_medium=distribute.pc_relevant.none-task-blog-2_defaultbaidujs_title~default-0.base&spm=1001.2101.3001.4242