File size: 1,092 Bytes
4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
def create_dataloaders(data_dir, batch_size=8, num_workers=2):
transform = transforms.Compose([
transforms.Resize((512, 512)), # keep 512p
transforms.Grayscale(num_output_channels=3), # convert to 3-ch for pre-trained ViT
transforms.ToTensor() # handle scaling
# Remove Normalize since doing reconstruction, Autoencoder -> target to be in same range as output (0,1)
#transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Std ViT normalization (ImageNet)
])
# Train only on 'normal' images
train_dataset = datasets.ImageFolder(root=f"{data_dir}/train", transform=transform)
test_dataset = datasets.ImageFolder(root=f"{data_dir}/test", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
return train_loader, test_loader |