def setup_dataset(args, train=True, val=True, test=True):
temporal_stride_size = args.temporal_stride
sample_duration = args.sample_duration * temporal_stride_size
normalization = make_normalization(args)
train_spatial_transform = [
Compose([
Scale(args.sample_size * 8 // 7),
MultiScaleCrop((args.sample_size, args.sample_size), args.scales),
# PhotometricDistort(),
ToTensor(args.norm_value),
normalization,
])
]
if args.hflip:
train_spatial_transform[0].transforms.insert(1, RandomHorizontalFlip())
temporal_stride = TemporalStride(temporal_stride_size)
train_temporal_transform = Compose([
temporal_stride,
TemporalRandomCrop(sample_duration / temporal_stride_size)
])
train_target_transform = ClassLabel()
train_data = make_dataset(args, 'training', train_spatial_transform,
train_temporal_transform,
train_target_transform) if train else None
# validation data
val_spatial_transform = [
Compose([
Scale(args.sample_size),
CenterCrop(args.sample_size),
ToTensor(args.norm_value), normalization
])
]
val_temporal_transform = Compose(
[temporal_stride,
LoopPadding(sample_duration / temporal_stride_size)])
val_target_transform = ClassLabel()
val_data = make_dataset(args, 'validation', val_spatial_transform,
val_temporal_transform,
val_target_transform) if val else None
# test data
test_spatial_transform = [
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
CornerCrop(args.sample_size, args.crop_position_in_test),
ToTensor(args.norm_value), normalization
])
]
if args.tta:
test_spatial_transform = []
for i in range(5):
test_spatial_transform.append(
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
GaussCrop(args.sample_size),
ToTensor(args.norm_value), normalization
]))
test_spatial_transform.append(
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
GaussCrop(args.sample_size),
HorizontalFlip(),
ToTensor(args.norm_value), normalization
]))
test_temporal_transform = Compose(
[temporal_stride,
LoopPadding(sample_duration / temporal_stride_size)])
test_target_transform = ClassLabel()
test_data = make_dataset(args, 'testing', test_spatial_transform,
test_temporal_transform,
test_target_transform) if test else None
return train_data, val_data, test_data
def setup_dataset(args, train=True, val=True, test=True):
temporal_stride_size = args.temporal_stride
sample_duration = args.sample_duration * temporal_stride_size
normalization = make_normalization(args)
photometric = []
if args.photometric:
photometric = [
# RandomContrast(),
RandomSharpness(lower=0.1),
RandomBrightness(delta=0.75),
# PhotometricDistort(),
]
train_spatial_transform = [
Compose([
Scale(args.sample_size),
RandomScale(scale_range=args.scales),
PadIfNeeded((args.sample_size, args.sample_size)),
MultiScaleCrop(
(args.sample_size,
args.sample_size), scale_ratios=[1.0]) if args.crop == 'fixed'
else RandomCrop(args.sample_size, mode=args.crop),
*photometric,
ToTensor(args.norm_value),
normalization,
])
]
if args.hflip:
train_spatial_transform[0].transforms.insert(1, RandomHorizontalFlip())
if args.vflip:
train_spatial_transform[0].transforms.insert(1, RandomVerticalFlip())
temporal_stride = TemporalStride(temporal_stride_size)
train_temporal_transform = Compose([
temporal_stride,
TemporalRandomCrop(sample_duration / temporal_stride_size)
])
train_target_transform = ClassLabel()
train_data = make_dataset(args, 'training', train_spatial_transform,
train_temporal_transform,
train_target_transform) if train else None
# validation data
val_spatial_transform = [
Compose([
Scale(args.sample_size),
CenterCrop(args.sample_size),
ToTensor(args.norm_value), normalization
])
]
val_temporal_transform = Compose(
[temporal_stride,
LoopPadding(sample_duration / temporal_stride_size)])
val_target_transform = ClassLabel()
val_data = make_dataset(args, 'validation', val_spatial_transform,
val_temporal_transform,
val_target_transform) if val else None
# test data
test_spatial_transform = [
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
CornerCrop(args.sample_size, args.crop_position_in_test),
ToTensor(args.norm_value), normalization
])
]
if args.tta:
test_spatial_transform = []
for i in range(5):
test_spatial_transform.append(
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
GaussCrop(args.sample_size),
ToTensor(args.norm_value), normalization
]))
test_spatial_transform.append(
Compose([
Scale(int(args.sample_size / args.scale_in_test)),
GaussCrop(args.sample_size),
HorizontalFlip(),
ToTensor(args.norm_value), normalization
]))
test_temporal_transform = Compose(
[temporal_stride,
LoopPadding(sample_duration / temporal_stride_size)])
test_target_transform = ClassLabel()
test_data = make_dataset(args, 'testing', test_spatial_transform,
test_temporal_transform,
test_target_transform) if test else None
return train_data, val_data, test_data