File size: 1,659 Bytes
3589275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import os
import torch
import torchvision
import torchvision.datasets as datasets
def rotate_img(img):
return torchvision.transforms.functional.rotate(img, -90)
def flip_img(img):
return torchvision.transforms.functional.hflip(img)
def emnist_preprocess():
return torchvision.transforms.Compose(
[
rotate_img,
flip_img,
]
)
class EMNIST:
def __init__(
self,
preprocess,
location,
batch_size=128,
num_workers=8,
):
preprocess1 = emnist_preprocess()
preprocess = torchvision.transforms.Compose(
[
preprocess,
preprocess1,
]
)
# if not os.path.exists(location):
# os.makedirs(location, exist_ok=True)
self.train_dataset = datasets.EMNIST(
root=location,
download=True,
split="digits",
transform=preprocess,
train=True,
)
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
)
self.test_dataset = datasets.EMNIST(
root=location,
download=True,
split="digits",
transform=preprocess,
train=False,
)
self.test_loader = torch.utils.data.DataLoader(
self.test_dataset,
batch_size=32,
shuffle=False,
num_workers=num_workers,
)
self.classnames = self.train_dataset.classes
|