【貓狗數據集】pytorch訓練貓狗數據集之創建數據集

數據集下載地址：

鏈接：https://pan.baidu.com/s/1l1AnBgkAAEhh0vI5_loWKw
提取碼：2xq4

貓狗數據集的分為訓練集25000張，在訓練集中貓和狗的圖像是混在一起的，pytorch讀取數據集有兩種方式，第一種方式是將不同類別的圖片放於其對應的類文件夾中，另一種是實現讀取數據集類，該類繼承torch.utils.Dataset，並重寫__getitem__和__len__。

 

先將貓和狗從訓練集中區分開來，分別放到dog和cat文件夾下：

import glob
import shutil
import os

#數據集目錄
path = "./ml/dogs-vs-cats/train"
#訓練集目錄
train_path = path+'/train'
#測試集目錄
test_path = path+'/test'

#將某類圖片移動到該類的文件夾下
def img_to_file(path):
    print("=========開始移動圖片============")
    #如果沒有dog類和cat類文件夾，則新建
    if not os.path.exists(path+"/dog"):
            os.makedirs(path+"/dog")
    if not os.path.exists(path+"/cat"):
            os.makedirs(path+"/cat")
    print("共：{}張圖片".format(len(glob.glob(path+"/*.jpg"))))
    #通過glob遍歷到所有的.jpg文件
    for imgPath in glob.glob(path+"/*.jpg"):
        #print(imgPath)
        #使用/划分
        img=imgPath.strip("\n").replace("\\","/").split("/")
        #print(img)
        #將圖片移動到指定的文件夾中
        if img[-1].split(".")[0] == "cat":
            shutil.move(imgPath,path+"/cat")
        if img[-1].split(".")[0] == "dog":
            shutil.move(imgPath,path+"/dog")
    print("=========移動圖片完成============")    
img_to_file(train_path)
print("訓練集貓共：{}張圖片".format(len(glob.glob(train_path+"/cat/*.jpg"))))
print("訓練集狗共：{}張圖片".format(len(glob.glob(train_path+"/dog/*.jpg"))))

然后從dog中和cat中分別抽取1250張，共2500張圖片作為測試集。

import random

def split_train_test(fileDir,tarDir):

        if not os.path.exists(tarDir):
            os.makedirs(tarDir)
        pathDir = os.listdir(fileDir)    #取圖片的原始路徑
        filenumber=len(pathDir)
        rate=0.1    #自定義抽取圖片的比例，比方說100張抽10張，那就是0.1
        picknumber=int(filenumber*rate) #按照rate比例從文件夾中取一定數量圖片
        sample = random.sample(pathDir, picknumber)  #隨機選取picknumber數量的樣本圖片
        print("=========開始移動圖片============")
        for name in sample:
                shutil.move(fileDir+name, tarDir+name)
        print("=========移動圖片完成============")
split_train_test(train_path+'/dog/',test_path+'/dog/')  
split_train_test(train_path+'/cat/',test_path+'/cat/')  

最終，我們就有以下結構了：

其中train包含22500張圖片，其中dog類和cat類各11250張。test包含2500張圖片，其中dog類和cat類各1250張。

發現測試集還是有點少，那就再來一遍了。

最后，train包含20250張圖片，其中dog類和cat類各10125張。test包含4750張圖片，其中dog類和cat類各2375張。