Skip to content

Instantly share code, notes, and snippets.

@MartinWeiss12
Created September 26, 2024 17:35
Show Gist options
  • Save MartinWeiss12/ce1d948a9d5f482f798010f2431ac20b to your computer and use it in GitHub Desktop.
Save MartinWeiss12/ce1d948a9d5f482f798010f2431ac20b to your computer and use it in GitHub Desktop.
Image Prep
image_width, image_height = 600, 450
data_transforms = {
'train': transforms.Compose([
transforms.Resize((image_width, image_height)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]),
'test': transforms.Compose([
transforms.Resize((image_width, image_height)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]),
}
# directory where all the train and test folders are
data_dir = 'mel-nv'
image_datasets = {
'train': datasets.ImageFolder(root=f'{data_dir}/train', transform=data_transforms['train']),
'test': datasets.ImageFolder(root=f'{data_dir}/test', transform=data_transforms['test'])
}
# DataLoader for batching
batch_size = 8
dataloaders = {
'train': DataLoader(image_datasets['train'], batch_size=batch_size, shuffle=True),
'test': DataLoader(image_datasets['test'], batch_size=batch_size, shuffle=False)
}
class_names = image_datasets['train'].classes
print(class_names)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment