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