麻雀牌を二段階で画像分類する｜4分類→牌種類分類のDatasetを作る

from PIL import Image
from pathlib import Path
from torch.utils.data import Dataset


class MahjongDataset(Dataset):
    GROUPS = {
        "manzu": list(range(0, 10)),
        "pinzu": list(range(10, 20)),
        "souzu": list(range(20, 30)),
        "jihai": list(range(30, 37)),
    }

    CLASS4_NAMES = [
        "manzu",
        "pinzu",
        "souzu",
        "jihai",
    ]

    def __init__(
        self,
        root_dir,
        split="train",
        task="class4",
        group=None,
        transform=None,
    ):
        """
        Parameters
        ----------
        root_dir
            datasetのルート
        split
            train / val / test
        task
            class4 または stage2
        group
            stage2時のみ
            manzu / pinzu / souzu / jihai
        """

        self.transform = transform
        self.samples = []

        root_dir = Path(root_dir) / split

        if not root_dir.exists():
            raise FileNotFoundError(root_dir)

        # -----------------------------
        # class4
        # -----------------------------
        if task == "class4":
            self.class_names = self.CLASS4_NAMES
            for class_index, group_name in enumerate(self.CLASS4_NAMES):
                tile_ids = self.GROUPS[group_name]
                for tile_id in tile_ids:
                    image_dir = root_dir / str(tile_id)
                    if not image_dir.exists():
                        continue
                    for image_path in image_dir.iterdir():
                        if image_path.suffix.lower() not in (".jpg"):
                            continue

                        self.samples.append((image_path, class_index,))

        # -----------------------------
        # stage2
        # -----------------------------
        elif task == "stage2":
            if group is None:
                raise ValueError("group must be specified.")

            if group not in self.GROUPS:
                raise ValueError(group)

            tile_ids = self.GROUPS[group]
            # モデルから返したいクラス名
            self.class_names = [str(i) for i in tile_ids]

            label_map = {
                tile_id: index
                for index, tile_id
                in enumerate(tile_ids)
            }

            for tile_id in tile_ids:
                image_dir = root_dir / str(tile_id)
                if not image_dir.exists():
                    continue
                for image_path in image_dir.iterdir():
                    if image_path.suffix.lower() not in (".jpg"):
                        continue

                    self.samples.append((image_path, label_map[tile_id],))

        else:
            raise ValueError(task)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        image_path, label = self.samples[index]
        image = Image.open(image_path).convert("RGB")
        if self.transform is not None:
            image = self.transform(image)
        return image, label