Pytorch – 自作のデータセットを扱う Dataset クラスを作る方法

from pathlib import Path

from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image


class ImageFolder(Dataset):
    IMG_EXTENSIONS = [".jpg", ".jpeg", ".png", ".bmp"]

    def __init__(self, img_dir, transform=None):
        # 画像ファイルのパス一覧を取得する。
        self.img_paths = self._get_img_paths(img_dir)
        self.transform = transform

    def __getitem__(self, index):
        path = self.img_paths[index]

        # 画像を読み込む。
        img = Image.open(path)

        if self.transform is not None:
            # 前処理がある場合は行う。
            img = self.transform(img)

        return img

    def _get_img_paths(self, img_dir):
        """指定したディレクトリ内の画像ファイルのパス一覧を取得する。
        """
        img_dir = Path(img_dir)
        img_paths = [
            p for p in img_dir.iterdir() if p.suffix in ImageFolder.IMG_EXTENSIONS
        ]

        return img_paths

    def __len__(self):
        """ディレクトリ内の画像ファイルの数を返す。
        """
        return len(self.img_paths)


# Transform を作成する。
transform = transforms.Compose([transforms.Resize(256), transforms.ToTensor()])
# Dataset を作成する。
dataset = ImageFolder("data", transform)
# DataLoader を作成する。
dataloader = DataLoader(dataset, batch_size=3)

for batch in dataloader:
    print(batch.shape)

from pathlib import Path

import cv2
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms


class Video(Dataset):
    def __init__(self, video_path, transform=None):
        self.cap = cv2.VideoCapture(video_path)
        self.transform = transform

    def __getitem__(self, index):
        """動画のフレームを返す。
        """
        # フレームを読み込む。
        ret, img = self.cap.read()
        # チャンネルの順番を変更する。 (B, G, R) -> (R, G, B)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # numpy 配列を PIL Image に変換する。
        img = Image.fromarray(img)

        if self.transform is not None:
            # 前処理がある場合は行う。
            img = self.transform(img)

        return img

    def __len__(self):
        """動画のフレーム数を返す。
        """
        return int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))


# Transform を作成する。
transform = transforms.Compose([transforms.Resize(256), transforms.ToTensor()])
# Dataset を作成する。
dataset = Video("data/sample.avi", transform)
# DataLoader を作成する。
dataloader = DataLoader(dataset, batch_size=128)

for batch in dataloader:
    print(batch.shape)

import pandas as pd
import torch
from sklearn.preprocessing import normalize
from torch.utils.data import DataLoader, Dataset


class Wine(Dataset):
    def __init__(self, csv_path):
        # csv ファイルを読み込む。
        df = pd.read_csv(csv_path)
        data = df.iloc[:, 1:]  # データ (2 ~ 14列目)
        labels = df.iloc[:, 0]  # ラベル (1列目)
        # データを標準化する。
        data = normalize(data)
        # クラス ID を 0 始まりにする。[1, 2, 3] -> [0, 1, 2]
        labels -= 1

        self.data = data
        self.labels = labels

    def __getitem__(self, index):
        """サンプルを返す。
        """
        return self.data[index], self.labels[index]

    def __len__(self):
        """csv の行数を返す。
        """
        return len(self.data)


# Dataset を作成する。
dataset = Wine("https://git.io/JfodD")
# DataLoader を作成する。
dataloader = DataLoader(dataset, batch_size=64)

for X_batch, y_batch in dataloader:
    print(X_batch.shape, y_batch.shape)

from pathlib import Path

from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.datasets import ImageFolder

# Transform を作成する。
transform = transforms.Compose([transforms.Resize(256), transforms.ToTensor()])
# Dataset を作成する。
dataset = ImageFolder("dataset1", transform)
print(dataset.class_to_idx)

# DataLoader を作成する。
dataloader = DataLoader(dataset, batch_size=3)

for X_batch, y_batch in dataloader:
    print(X_batch.shape, y_batch.shape)