def create_loaders(train_dir, val_dir, train_list, val_list, shorter_side, crop_size, low_scale, high_scale, normalise_params, batch_size, num_workers, ignore_label): """ Args: train_dir (str) : path to the root directory of the training set. val_dir (str) : path to the root directory of the validation set. train_list (str) : path to the training list. val_list (str) : path to the validation list. shorter_side (int) : parameter of the shorter_side resize transformation. crop_size (int) : square crop to apply during the training. low_scale (float) : lowest scale ratio for augmentations. high_scale (float) : highest scale ratio for augmentations. normalise_params (list / tuple) : img_scale, img_mean, img_std. batch_size (int) : training batch size. num_workers (int) : number of workers to parallelise data loading operations. ignore_label (int) : label to pad segmentation masks with Returns: train_loader, val loader """ # Torch libraries from torchvision import transforms from torch.utils.data import DataLoader, random_split # Custom libraries if args.data_name == 'nyu': from datasets import NYUDataset as Dataset elif args.data_name == 'cityscapes': from datasets import CSDataset as Dataset from datasets import Pad, RandomCrop, RandomMirror, ResizeShorterScale, ToTensor, Normalise ## Transformations during training ## composed_trn = transforms.Compose([ ResizeShorterScale(shorter_side, low_scale, high_scale), Pad(crop_size, [123.675, 116.28, 103.53], ignore_label), RandomMirror(), RandomCrop(crop_size), Normalise(*normalise_params), ToTensor() ]) composed_val = transforms.Compose( [Normalise(*normalise_params), ToTensor()]) ## Training and validation sets ## trainset = Dataset(data_file=train_list, data_dir=train_dir, transform_trn=composed_trn, transform_val=composed_val) valset = Dataset(data_file=val_list, data_dir=val_dir, transform_trn=None, transform_val=composed_val) logger.info( " Created train set = {} examples, val set = {} examples".format( len(trainset), len(valset))) ## Training and validation loaders ## train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True) val_loader = DataLoader(valset, batch_size=1, shuffle=False, num_workers=num_workers, pin_memory=True) return train_loader, val_loader
def create_loaders(train_dir, val_dir, train_list, val_list, shorter_side, crop_size, low_scale, high_scale, normalise_params, batch_size, num_workers, ignore_label): """ Args: train_dir (str) : path to the root directory of the training set. val_dir (str) : path to the root directory of the validation set. train_list (str) : path to the training list. val_list (str) : path to the validation list. shorter_side (int) : parameter of the shorter_side resize transformation. crop_size (int) : square crop to apply during the training. low_scale (float) : lowest scale ratio for augmentations. high_scale (float) : highest scale ratio for augmentations. normalise_params (list / tuple) : img_scale, img_mean, img_std, depth_scale, depth_meam, depth_std batch_size (int) : training batch size. num_workers (int) : number of workers to parallelise data loading operations. ignore_label (int) : label to pad segmentation masks with Returns: train_loader, val loader """ # Torch libraries from torchvision import transforms from torch.utils.data import DataLoader, random_split # Custom libraries from datasets import NYUDataset as Dataset from datasets import Pad, RandomCrop, RandomMirror, ResizeShorterScale, ToTensor, Normalise, RGBCutout ## Transformations during training ## ### modified to take HHA depth image as well composed_trn = transforms.Compose([ ResizeShorterScale(shorter_side, low_scale, high_scale), Pad(crop_size, [123.675, 116.28, 103.53], [111.0, 113.0, 133.0], ignore_label), RandomMirror(), RandomCrop(crop_size), # 165 cutout size is ~1/3 of the 500 size image. a guess based on the # paper's cifar-10 and cifar-100 values, given our 40 classes RGBCutout([123.675, 116.28, 103.53], 165), Normalise(*normalise_params), ToTensor() ]) composed_val = transforms.Compose( [Normalise(*normalise_params), ToTensor()]) ## Training and validation sets ## trainset = Dataset(data_file=train_list, data_dir=train_dir, transform_trn=composed_trn, transform_val=composed_val) valset = Dataset(data_file=val_list, data_dir=val_dir, transform_trn=None, transform_val=composed_val) logger.info( " Created train set = {} examples, val set = {} examples".format( len(trainset), len(valset))) ## Training and validation loaders ## train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True) val_loader = DataLoader(valset, batch_size=1, shuffle=False, num_workers=num_workers, pin_memory=True) return train_loader, val_loader