pytorch?dataset實戰(zhàn)案例之讀取數據集的代碼

更新時間：2022年10月21日 10:59:17 作者：半島鐵子_

這篇文章主要介紹了pytorch?dataset實戰(zhàn)案例讀取數據集的相關知識,這段代碼可以作為讀取數據集的一個DataSet類的基礎類，可以擴充進行修改，以后有類似需要可以拿過來修改，本文結合實例代碼給大家詳細講解，需要的朋友可以參考下

概述

最近在跑一篇圖像修復論文的代碼，配置好環(huán)境之后開始運行，發(fā)現(xiàn)數據一直加載不進去。
害，還是得看人家代碼咋寫的，一句一句看邏輯，準能找出問題。通讀dataset后，發(fā)現(xiàn)了問題所在，終于成功加載了數據集。

項目結構與代碼

項目結構

主要的目的就是從數據集中讀取到彩色圖像和掩碼圖像。
代碼
代碼中涉及到torch.transforms、合并路徑等知識點，我在代碼中都進行了詳細的注釋，路徑要對照著項目結構，如果自己用的話要根據項目結構去將相對路徑改過來。
dataset.py ：當前的工作路徑：…\OT-GAN-for-Inpainting-master\src\data

import os
import math
import numpy as np
from glob import glob

from random import shuffle
from PIL import Image, ImageFilter

import torch
import torchvision.transforms.functional as F
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader



class InpaintingData(Dataset):
    def __init__(self, args):
        super(Dataset, self).__init__()     # 繼承Dataset的父類的初始化函數
        self.w = self.h = args.image_size   # 通過args傳入新的屬性---圖像的w和h
        self.mask_type = args.mask_type     # 通過args傳入新的屬性---mask_type

        # image and mask 
        self.image_path = []    #創(chuàng)建image_path的數組
        for ext in ['*.jpg', '*.png']:  # 獲取每一個后綴為.jpg或者.png的圖片，為ext
            # 將dir_image、data_train和ext拼接作為圖片的路徑，并將其存入到數組image_path之中,glob()獲取一個lsit集合
            self.image_path.extend(glob(os.path.join(args.dir_image, args.data_train, ext)))
        self.mask_path = glob(os.path.join(args.dir_mask, args.mask_type, '*.png')) #拼接dir_mask、mask_type和路徑下所有的.png作為mask_path

        # augmentation 
        self.img_trans = transforms.Compose([   #接收一個 transforms方法的list為參數，將這些操作組合到一起，返回一個新的tranforms
            transforms.RandomResizedCrop(args.image_size),  #隨機隨機長寬比裁剪，大小為image_size
            transforms.RandomHorizontalFlip(), #隨機水平翻轉
            transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), #改變圖像的亮度、對比度、飽和度和色調。
            transforms.ToTensor()])     # 轉為tensor，并歸一化至[0-1]
        self.mask_trans = transforms.Compose([
            transforms.Resize(args.image_size, interpolation=transforms.InterpolationMode.NEAREST), #將輸入圖像調整為給定的大小，interpolation是插值方式，此處是默認值NEAREST
            transforms.RandomHorizontalFlip(),  #隨機水平翻轉
            transforms.RandomRotation(  #隨機旋轉
                (0, 45), interpolation=transforms.InterpolationMode.NEAREST), #(0, 45)是角度
        ])

    def __len__(self):  # __len__和__getitem__DataSet類必須實現(xiàn)的靜態(tài)方法
        return len(self.image_path)

    def __getitem__(self, index):
        # load image
        image = Image.open(self.image_path[index]).convert('RGB') #獲取圖像，并將其轉化為RGB（3x8位像素）模式
        filename = os.path.basename(self.image_path[index]) #獲取圖片的路徑

        if self.mask_type == 'pconv': #如果mask_type為pconv
            index = np.random.randint(0, len(self.mask_path)) #隨機從mask_path中獲取一個下標
            mask = Image.open(self.mask_path[index])    #根據下標獲取mask圖片
            mask = mask.convert('L')    #將mask圖片轉化為L(8位像素的黑白圖片，0表示黑，255表示白)模式
        else:   # 構造mask，有mask數據集的話就運行不到這里
            mask = np.zeros((self.h, self.w)).astype(np.uint8) #構造與h和w一樣大的圖片，都用0填充，并將其轉換為uint8
            mask[self.h // 4:self.h // 4 * 3, self.w // 4:self.w // 4 * 3] = 1
            mask = Image.fromarray(m).convert('L')

        # augment
        image = self.img_trans(image) * 2. - 1. # 數據標準化，將輸出限定在一定的范圍
        mask = F.to_tensor(self.mask_trans(mask))   # 將轉化后的mask圖像轉化為tensor

        return image, mask, filename    #返回


if __name__ == '__main__':
    from attrdict import AttrDict

    args = {
        'dir_image': '../../examples/logos',
        'data_train': 'image',
        'dir_mask': '../../examples/logos/mask',
        'mask_type': 'pconv',
        'image_size': 512
    }
    args = AttrDict(args) # 將上面定義的參數傳入AttrDict()作為新參數

    data = InpaintingData(args)     #創(chuàng)建InpaintingData對象
    print(len(data), len(data.mask_path))   #輸出data的長度，mask的長度
    img, mask, filename = data[0]   # 獲取第一張圖片
    print(img.size(), mask.size(), filename)    #打印上述信息

輸出：