| from typing import Callable |
|
|
| import numpy as np |
| from torchvision.datasets import MNIST |
|
|
| from ..config import Config |
| from ..dataset import BaseDataModule |
|
|
|
|
| class MNISTDataset(MNIST): |
| def __init__( |
| self, |
| train: bool = True, |
| preprocess: None | Callable = None, |
| augmentations: None | Callable = None, |
| ): |
| super().__init__(root="datasets/other/", train=train, download=True) |
| self.preprocess = preprocess |
| self.augmentations = augmentations |
|
|
| def __getitem__(self, idx): |
| image, label = super().__getitem__(idx) |
| if self.augmentations is not None: |
| image = self.augmentations(image) |
| if self.preprocess is not None: |
| image = self.preprocess(image) |
| return { |
| "image": image, |
| "label": label, |
| "path": f"{idx}: {label}", |
| "idx": idx, |
| } |
|
|
| def print_statistics(self): |
| print(f"Number of samples: {len(self)}") |
| unique, counts = np.unique(self.targets, return_counts=True) |
| print("Class distribution") |
| names = self.get_class_names() |
| for u, c in zip(unique, counts): |
| print(f"Class {u} ({names[u]}): {c}") |
|
|
| def get_class_names(self) -> dict[int, str]: |
| return {i: str(i) for i in range(10)} |
|
|
|
|
| class MNISTDataModule(BaseDataModule): |
| def __init__(self, config: Config, preprocess: None | Callable = None): |
| super().__init__(config, preprocess) |
|
|
| def setup(self, stage: str): |
| |
| if stage == "fit" or stage == "validate": |
| self.train_dataset = MNISTDataset(train=True, preprocess=self.preprocess) |
| self.val_dataset = MNISTDataset(train=False, preprocess=self.preprocess) |
|
|
| print("\nTrain dataset") |
| self.train_dataset.print_statistics() |
|
|
| print("\nValidation dataset") |
| self.val_dataset.print_statistics() |
|
|
| if stage == "test": |
| self.test_dataset = MNISTDataset(train=False, preprocess=self.preprocess) |
|
|
| print("\nTest dataset") |
| self.test_dataset.print_statistics() |
|
|
