Spaces:
Sleeping
Sleeping
| import os | |
| from torchvision import datasets, transforms | |
| from torch.utils.data import DataLoader | |
| def get_data_loaders(data_dir="data/chest_xray", batch_size=4): | |
| transform = transforms.Compose([ | |
| transforms.Resize((128,128)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.5],[0.5]) | |
| ]) | |
| train_dataset = datasets.ImageFolder( | |
| os.path.join(data_dir, "train"), transform=transform | |
| ) | |
| val_dataset = datasets.ImageFolder( | |
| os.path.join(data_dir, "val"), transform=transform | |
| ) | |
| test_dataset = datasets.ImageFolder( | |
| os.path.join(data_dir, "test"), transform=transform | |
| ) | |
| train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0) | |
| val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=0) | |
| test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=0) | |
| return train_loader, val_loader, test_loader |