本文將原始的numpy array數據在pytorch下封裝爲Dataset類的數據集，爲後續深度網絡訓練提供數據。
加載並保存圖像信息
首先導入需要的庫，定義各種路徑。
import os
import matplotlib
from keras.datasets import mnist
import numpy as np
from torch.utils.data.dataset import Dataset
from PIL import Image
import scipy.misc
root_path = 'E:/coding_ex/pytorch/Alexnet/data/'
base_path = 'baseset/'
training_path = 'trainingset/'
test_path = 'testset/'
這裏將數據集分爲三類，baseset爲所有數據（trainingset+testset），trainingset是訓練集，testset是測試集。
直接通過keras.dataset加載mnist數據集，不能自動下載的話可以手動下載.npz並保存至相應目錄下。
def LoadData(root_path, base_path, training_path, test_path):
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_baseset = np.concatenate((x_train, x_test))
y_baseset = np.concatenate((y_train, y_test))
train_num = len(x_train)
test_num = len(x_test)
#baseset
file_img = open((os.path.join(root_path, base_path)+'baseset_img.txt'),'w')
file_label = open((os.path.join(root_path, base_path)+'baseset_label.txt'),'w')
for i in range(train_num + test_num):
file_img.write(root_path + base_path + 'img/' + str(i) + '.png\n') #name
file_label.write(str(y_baseset[i])+'\n') #label
# scipy.misc.imsave(root_path + base_path + '/img/'+str(i) + '.png', x_baseset[i])
matplotlib.image.imsave(root_path + base_path + 'img/'+str(i) + '.png', x_baseset[i])
file_img.close()
file_label.close()
#trainingset
file_img = open((os.path.join(root_path, training_path)+'trainingset_img.txt'),'w')
file_label = open((os.path.join(root_path, training_path)+'trainingset_label.txt'),'w')
for i in range(train_num):
file_img.write(root_path + training_path + 'img/' + str(i) + '.png\n') #name
file_label.write(str(y_train[i])+'\n') #label
# scipy.misc.imsave(root_path + training_path + '/img/'+str(i) + '.png', x_train[i])
matplotlib.image.imsave(root_path + training_path + 'img/'+str(i) + '.png', x_train[i])
file_img.close()
file_label.close()
#testset
file_img = open((os.path.join(root_path, test_path)+'testset_img.txt'),'w')
file_label = open((os.path.join(root_path, test_path)+'testset_label.txt'),'w')
for i in range(test_num):
file_img.write(root_path + test_path + 'img/' + str(i) + '.png\n') #name
file_label.write(str(y_test[i])+'\n') #label
# scipy.misc.imsave(root_path + test_path + '/img/'+str(i) + '.png', x_test[i])
matplotlib.image.imsave(root_path + test_path + 'img/'+str(i) + '.png', x_test[i])
file_img.close()
file_label.close()
使用這段代碼時，需要建立相應的文件夾及.txt文件，./data文件夾結構如下：
/img文件夾
由於mnist數據集其實是灰度圖，這裏用matplotlib保存的圖像是僞彩色圖像。
如果用scipy.misc.imsave的話保存的則是灰度圖像。
xxx_img.txt文件
xxx_img.txt文件中存放的是每張圖像的名字
xxx_label.txt文件
xxx_label.txt文件中存放的是類別標記
這裏記得保存的時候一行爲一個圖像信息，便於後續讀取。
定義自己的Dataset類
pytorch訓練數據時需要數據集爲Dataset類，便於迭代等等，這裏將加載保存之後的數據封裝成Dataset類，繼承該類需要寫初始化方法（init），獲取指定下標數據的方法（getitem），獲取數據個數的方法（len）。
class DataProcessingMnist(Dataset):
def __init__(self, root_path, imgfile_path, labelfile_path, imgdata_path, transform = None):
self.root_path = root_path
self.transform = transform
self.imagedata_path = imgdata_path
img_file = open((root_path + imgfile_path),'r')
self.image_name = [x.strip() for x in img_file]#這裏要將回車換行符去掉
img_file.close()
label_file = open((root_path + labelfile_path), 'r')
self.label = [int(x.strip()) for x in label_file]#去除回車換行符
label_file.close()
def __getitem__(self, idx):
image = Image.open(str(self.image_name[idx]))
if self.transform is not None:
image = self.transform(image)
label = self.label[idx]
return image, label
def __len__(self):
return len(self.image_name)
定義完自己的類之後可以測試一下。
LoadData(root_path, base_path, training_path, test_path)
training_imgfile = training_path + 'trainingset_img.txt'
training_labelfile = training_path + 'trainingset_label.txt'
training_imgdata = training_path + 'img/'
#實例化一個類
dataset = DataProcessingMnist(root_path, training_imgfile, training_labelfile, training_imgdata)
得到圖像名稱
name = dataset.image_name
這裏我們可以單獨輸出某一個名稱看一下是否有換行符
print(name[0])
>>>'E:/coding_ex/pytorch/Alexnet/data/trainingset/img/0.png'
如果定義類的時候self.image_name = [x.strip() for x in img_file]這句沒有strip掉，則輸出的值將爲’E:/coding_ex/pytorch/Alexnet/data/trainingset/img/0.png\n’
獲取固定下標的圖像
im, label = dataset.__getitem__(0)
得到結果
