File size: 1,525 Bytes
28e6f98 | 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 | from torchvision import transforms
from torch.utils.data import DataLoader
def get_dataset(name, args):
batch_size = args.batch_size * len(args.gpu.split(','))
if name == 'brats':
from dataloaders.BRATS_dataloader_new import Hybrid as BratsDataset
from dataloaders.BRATS_dataloader_new import RandomPadCrop, ToTensor, RandomFlip
db_train = BratsDataset(split='train', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR,
transform=transforms.Compose([RandomPadCrop(), ToTensor()]), # RandomFlip(),
base_dir=args.root_path, input_normalize = args.input_normalize,
use_kspace=args.use_kspace)
db_test = BratsDataset(split='test', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR,
transform=transforms.Compose([ToTensor()]),
base_dir=args.root_path, input_normalize = args.input_normalize,
use_kspace=args.use_kspace)
trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
fixtrainloader = DataLoader(db_train, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
return trainloader, fixtrainloader, testloader
else:
raise NotImplementedError(f'Dataset {name} is not implemented.')
|