pytorch dataset實戰案例之讀取資料集的程式碼

2022-10-23 18:01:16

概述

最近在跑一篇影象修復論文的程式碼，設定好環境之後開始執行，發現資料一直載入不進去。
害，還是得看人家程式碼咋寫的，一句一句看邏輯，準能找出問題。通讀dataset後，發現了問題所在，終於成功載入了資料集。

專案結構與程式碼

專案結構

主要的目的就是從資料集中讀取到彩色影象和掩碼影象。
程式碼
程式碼中涉及到torch.transforms、合併路徑等知識點，我在程式碼中都進行了詳細的註釋，路徑要對照著專案結構，如果自己用的話要根據專案結構去將相對路徑改過來。
dataset.py ：當前的工作路徑：…OT-GAN-for-Inpainting-mastersrcdata

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 = []    #建立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類必須實現的靜態方法
        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)     #建立InpaintingData物件
    print(len(data), len(data.mask_path))   #輸出data的長度，mask的長度
    img, mask, filename = data[0]   # 獲取第一張圖片
    print(img.size(), mask.size(), filename)    #列印上述資訊

輸出：