* File: __init__.py

* Created on: Tue Sep 14 20:39:57 2021

* Author: Liam Woolley

* Licence: MIT

"""

import torch

from torch.utils.data import Dataset, DataLoader

class MNISTDataset(Dataset):

def __init__(self, image_dir='./images', label_dir='./labels'):

self.image_dir = image_dir

self.label_dir = label_dir

super().__init__()

def __getitem__(self, index):

img = torch.load(f'{self.image_dir}/{index}.jpg')

return (img, torch.tensor(int(f'{self.label_dir}/{index}.txt'))), None

def __len__(self):

return len(os.listdir(self.image_dir))

class MNISTDataLoader(DataLoader):

def __init__(self, dataset, batch_size=32, shuffle=True, num_workers=4):

super().__init__(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

self.dataset = dataset

def __getitem__(self, index):

img, label = self.dataset.__getitem__(index)

return torch.tensor(img), label

def load_data():

ds = MNISTDataset()

dl = MNISTDataLoader(ds)

return ds, dl