Spaces:
Sleeping
Sleeping
File size: 2,320 Bytes
2a2559c 106788a 2a2559c | 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 | import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms
from .config import DataPreprocessorConfigsInput
class DataPreprocessor:
def __init__(self, configs: DataPreprocessorConfigsInput):
self.configs = configs
self.train_dataset_path = self.configs.train_dataset_path
self.test_dataset_path = self.configs.test_dataset_path
self.shuffle = self.configs.shuffle
self.batch_size = self.configs.batch_size
self.horizontal_flip_prob = self.configs.horizontal_flip_prob # Probability for random horizontal flip
self.image_size = self.configs.image_size
self.mean = self.configs.mean
self.std = self.configs.std
def preprocess(self, label:str = "train") -> transforms.Compose:
if label == "train":
transform = transforms.Compose([
transforms.Resize((self.image_size, self.image_size)),
transforms.RandomHorizontalFlip(p=self.horizontal_flip_prob),
transforms.RandomRotation(degrees=15),
transforms.ToTensor(),
transforms.Normalize(mean=self.mean, std=self.std),
])
else:
transform = transforms.Compose([
transforms.Resize((self.image_size, self.image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=self.mean, std=self.std)
])
return transform
def create_dataloader(self) -> DataLoader:
train_transform = self.preprocess()
test_transform = self.preprocess(label="test")
train_dataset = datasets.ImageFolder(root=self.train_dataset_path, transform=train_transform)
train_dataloader = DataLoader(
dataset=train_dataset,
batch_size=self.batch_size,
shuffle=self.shuffle
)
test_dataset = datasets.ImageFolder(root=self.test_dataset_path, transform=test_transform)
test_dataloader = DataLoader(
dataset=test_dataset,
batch_size=self.batch_size,
shuffle= not self.shuffle
)
print(f"Train dataset size: {len(train_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")
return train_dataloader, test_dataloader |