def __call__(self, image, annot):
if self._RAND_CROP:
image, annot = self._rand_crop(image, annot)
if self._NORMAL_SCALE:
scale_rate = (random.random() *
(self._scale_range[1] - self._scale_range[0]) +
self._scale_range[0])
image = cv2.resize(image, None, fx=scale_rate, fy=scale_rate)
annot = cv2.resize(
annot,
None,
fx=scale_rate,
fy=scale_rate,
interpolation=cv2.INTER_NEAREST,
)
if self._F_RAND_FLAG:
self._transformer.rand_f(self._F_RANGE)
if self._EXT_RAND_FLAG:
self._transformer.rand_ext_params()
dst_image = self._transformer.transFromColor(image)
dst_annot = self._transformer.transFromGray(annot, reuse=True)
if self._NORMAL_TRANSLATE:
x_shift = (random.random() *
(self._trans_range[1] - self._trans_range[0]) +
self._trans_range[0])
y_shift = (random.random() *
(self._trans_range[1] - self._trans_range[0]) +
self._trans_range[0])
M = np.array([[1, 0, x_shift], [0, 1, y_shift]], dtype=np.float32)
sz = (dst_annot.shape[1], dst_annot.shape[0])
dst_image = cv2.warpAffine(dst_image, M, sz)
dst_annot = cv2.warpAffine(dst_annot,
M,
sz,
flags=cv2.INTER_NEAREST,
borderValue=20)
if self._FISH_SCALE:
dst_image, dst_annot = self._fish_scale(dst_image, dst_annot)
dst_image = Image.fromarray(dst_image)
dst_annot = Image.fromarray(dst_annot)
brightness, contrast, hue, saturation = 0.1, 0.1, 0.1, 0.1
dst_image = ColorJitter(brightness, contrast, saturation,
hue)(dst_image)
if random.random() < 0.5:
dst_image = dst_image.transpose(Image.FLIP_LEFT_RIGHT)
dst_annot = dst_annot.transpose(Image.FLIP_LEFT_RIGHT)
dst_annot = np.asarray(dst_annot)
dst_annot = torch.from_numpy(dst_annot)
dst_image = ToTensor()(dst_image)
# image = Normalize(mean=[.485, .456, .406], std=[.229, .224, .225])(image)
return dst_image, dst_annot
def get_recognition_transform(resolution):
return Compose([
RandomResizedCrop(size=(resolution, resolution), scale=(.5, 1.)),
RandomHorizontalFlip(),
RandomVerticalFlip(p=.05),
ColorJitter(brightness=.25),
ColorJitter(contrast=.25),
ColorJitter(saturation=.25),
ColorJitter(hue=.25),
Resize((resolution, resolution)),
ToTensor()
])
def __init__(self, data, info=None):
assert data['observations'].dtype == np.uint8
self.size = data['observations'].shape[0]
self.traj_length = data['observations'].shape[1]
self.data = data
self.info = info
self.jitter = ColorJitter((0.75,1.25), (0.9,1.1), (0.9,1.1), (-0.1,0.1))
#self.crop = RandomResizedCrop((48, 48), (0.9, 0.9), (1, 1))
# RandomResizedCrop((int(sqrt(self.imlength)), int(sqrt(self.imlength))), (0.9, 0.9), (1, 1))
if 'env' not in self.data:
self.data['env'] = self.data['observations'][:, 0, :]
def __call__(self, sample):
image, label = sample['image'], sample['label']
if random.random() < self.brightness[0]:
pytorch_color_jitter = PytorchColorJitter(
brightness=self.brightness[1], contrast=0, saturation=0, hue=0)
image = pytorch_color_jitter.forward(image)
if random.random() < self.contrast[0]:
pytorch_color_jitter = PytorchColorJitter(
brightness=0, contrast=self.contrast[1], saturation=0, hue=0)
image = pytorch_color_jitter.forward(image)
if random.random() < self.saturation[0]:
pytorch_color_jitter = PytorchColorJitter(
brightness=0, contrast=0, saturation=self.saturation[1], hue=0)
image = pytorch_color_jitter.forward(image)
if random.random() < self.hue[0]:
pytorch_color_jitter = PytorchColorJitter(brightness=0,
contrast=0,
saturation=0,
hue=self.hue[1])
image = pytorch_color_jitter.forward(image)
return {'image': image, 'label': label}
def test_single_transform(self):
from torchvision.transforms import Compose, RandomCrop, RandomRotation, ColorJitter, ToTensor, RandomResizedCrop
transforms = SequentialWrapper(com_transform=Compose([
RandomRotation(45),
RandomCrop(224),
RandomResizedCrop(size=192, scale=(0.8, 1.2))
]),
image_transform=Compose([
ColorJitter(brightness=[0.8, 1.2],
contrast=[0.8, 1.2],
saturation=1),
ToTensor()
]),
target_transform=ToLabel())
dataset = ACDCDataset(
root_dir=self._root,
mode="train",
transforms=transforms,
)
(image, target), filename = dataset[4]
from deepclustering3.viewer import multi_slice_viewer_debug
import matplotlib.pyplot as plt
multi_slice_viewer_debug(image, target, no_contour=True)
plt.show()
def __init__(self, datadir, resize, argument=True):
self.input_images_dir = os.path.join(datadir, 'images')
self.masks_dir = os.path.join(datadir, 'masks')
self.argument = argument
self.pairlist = self.search_pairs()
if self.argument == False:
self.image_transforms = Compose([
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='celeba')
if resize > 0:
self.joint_transforms = ComposeX([
Resize(resize),
])
else:
self.joint_transforms = None
else:
self.image_transforms = Compose([
ColorJitter(0.05, 0.05, 0.05),
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='celeba')
if resize > 0:
self.joint_transforms = ComposeX(
[Resize(resize), RandomHorizontallyFlip()])
else:
self.joint_transforms = RandomHorizontallyFlip()
def __init__(self,
base_dataset,
img_size=(32, 32),
color_dist_strength=0.5):
"""
Parameters
----------
base_dataset : Dataset
img_size : tuple
color_dist_strength : float
"""
self.base_dataset = base_dataset
# get color distortion random transform, Color distortion (see Pytorch
# pseudocode in arXiv:2002.05709)
color_jitter = ColorJitter(
0.8 * color_dist_strength,
0.8 * color_dist_strength,
0.8 * color_dist_strength,
0.2 * color_dist_strength,
)
rnd_color_jitter = RandomApply([color_jitter], p=0.8)
rnd_gray = RandomGrayscale(p=0.2)
self.transform = Compose([
RandomResizedCrop(img_size), rnd_color_jitter, rnd_gray,
ToTensor()
])
def get_transform(
target_size=(288, 288),
transform_list='random_crop, horizontal_flip', # random_crop | keep_aspect
augment_ratio=0.5,
is_train=True,
):
transform = list()
transform_list = transform_list.split(', ')
augments = list()
for transform_name in transform_list:
if transform_name == 'random_crop':
scale = (0.5, 1.0) if is_train else (0.8, 1.0)
transform.append(RandomResizedCrop(target_size, scale=(0.8, 1.0)))
elif transform_name == 'keep_aspect':
transform.append(KeepAsepctResize(target_size))
elif transform_name == 'horizontal_flip':
augments.append(RandomHorizontalFlip())
elif transform_name == 'vertical_flip':
augments.append(RandomVerticalFlip())
elif transform_name == 'random_rotate':
augments.append(RandomRotate())
elif transform_name == 'color_jitter':
brightness = 0.1 if is_train else 0.05
contrast = 0.1 if is_train else 0.05
augments.append(
ColorJitter(
brightness=brightness,
contrast=contrast,
saturation=0,
hue=0,
))
transform.append(RandomApply(augments, p=augment_ratio))
return Compose(transform)
def get_data_loaders(train_files, val_files, img_size=224):
train_transform = Compose([
ColorJitter(0.3, 0.3, 0.3, 0.3),
RandomResizedCrop(img_size, scale=(0.8, 1.2)),
RandomAffine(10.),
RandomRotation(13.),
RandomHorizontalFlip(),
ToTensor(),
])
#train_mask_transform = Compose([
# RandomResizedCrop(img_size, scale=(0.8, 1.2)),
# RandomAffine(10.),
# RandomRotation(13.),
# RandomHorizontalFlip(),
# ToTensor(),
#])
val_transform = Compose([
Resize((img_size, img_size)),
ToTensor(),
])
train_loader = DataLoader(MaskDataset(train_files, train_transform),
batch_size=BATCH_SIZE,
shuffle=True,
pin_memory=True,
num_workers=4)
val_loader = DataLoader(MaskDataset(val_files, val_transform),
batch_size=BATCH_SIZE,
shuffle=False,
pin_memory=True,
num_workers=4)
return train_loader, val_loader
def get_segmentation_transform(self) -> ModelTransformsPerExecutionMode:
if self.imaging_feature_type in [ImagingFeatureType.Segmentation, ImagingFeatureType.ImageAndSegmentation]:
return ModelTransformsPerExecutionMode(
train=ImageTransformationPipeline(
transforms=[RandomAffine(10), ColorJitter(0.2)],
use_different_transformation_per_channel=True))
return ModelTransformsPerExecutionMode()
def SSDTransform(size,
color_jitter=True,
scale=(0.1, 1),
expand=(1, 4),
min_area_frac=0.25):
transforms = []
if color_jitter:
transforms.append(
InputTransform(
ColorJitter(
brightness=0.1,
contrast=0.5,
saturation=0.5,
hue=0.05,
)))
transforms += [
RandomApply([
RandomExpand(expand),
]),
RandomChoice([
UseOriginal(),
RandomSampleCrop(),
RandomResizedCrop(size,
scale=scale,
ratio=(1 / 2, 2 / 1),
min_area_frac=min_area_frac),
]),
RandomHorizontalFlip(),
Resize(size)
]
return Compose(transforms)
def _get_img_transformations(means: typing.List[float],
stds: typing.List[float]
) -> Compose:
""" Get all transformations to carry out on image.
Helpful while creating GestureDataset.
Arguments
----------
means : typing.Tuple
Mean values for data normalization
stds : typing.Tuple
Std values for data normalization
Returns
-------
torchvision.transforms.Compose
composition of transformations
"""
return Compose([
Resize((512, 256)),
ColorJitter(brightness=0.1, contrast=0.1), # will be applied RANDOMLY
RandomVerticalFlip(p=0.1),
ToTensor(),
Normalize(means, stds),
RandomApply([Lambda(lambda x: x + torch.randn_like(x))], p=0.3) # noise
])
def __init__(self, data_dir, out_size=800):
if 'train' in data_dir:
self.bg_dir = data_dir[0:data_dir.find('train')] + 'background'
else:
self.bg_dir = data_dir[0:data_dir.find('test')] + 'background'
self.bg_name = [
name for name in os.listdir(self.bg_dir)
if os.path.isfile(os.path.join(self.bg_dir, name))
]
self.num_bg = len(self.bg_name)
self.bg_object = [
open_img(os.path.join(self.bg_dir, bg_name)).convert('RGB')
for bg_name in self.bg_name
]
self.bg_scale = (0.5, 1.0)
self.obj_scale = (0.2, 0.4, 0.6, 0.8, 1)
#self.obj_size = obj_size
self.bg_size = out_size
self.rrc = RandomResizedCrop(size=out_size, scale=self.bg_scale)
self.rrc_obj = self.class_resize(min_size=500,
max_size=(600, 575, 550))
self.color_jig = ColorJitter()
def __getitem__(self, i):
record = self.selection.iloc[i]
imgpath = os.path.join(self.rootdir, record['imgpath'])
_img = Image.open(imgpath)
if self.target_aspect[i] < 1:
height = self.allowed_size
width = int(height * self.target_aspect[i])
else:
width = self.allowed_size
height = int(width / self.target_aspect[i])
_img = _img.convert('RGB').resize((width, height), Image.BILINEAR)
_resizer = RandomResizedCrop(min(width, height))
_jitterer = ColorJitter(0.1, 0.1, 0.1, 0.05)
_img = _jitterer(_resizer(_img))
img = T.FloatTensor(np.array(_img)).permute(2, 0, 1) / 255.
_img.close()
return img, \
T.LongTensor([record.cat_id]), \
T.FloatTensor([[(record.xmin + record.xmax) / 2,
(record.ymin + record.ymax) / 2,
(record.xmax - record.xmin),
(record.ymax - record.ymin),]])
def test_twice_transform(self):
from torchvision.transforms import Compose, RandomCrop, RandomRotation, ColorJitter, ToTensor
transforms = SequentialWrapperTwice(com_transform=Compose(
[RandomRotation(45), RandomCrop(224)], ),
image_transform=Compose([
ColorJitter(
brightness=[0.8, 1.2],
contrast=[0.8, 1.2],
saturation=1),
ToTensor()
]),
target_transform=ToLabel(),
total_freedom=False)
dataset = ACDCDataset(
root_dir=self._root,
mode="train",
transforms=transforms,
)
(image1, image2, target1, target2), filename = dataset[4]
from deepclustering3.viewer import multi_slice_viewer_debug
import matplotlib.pyplot as plt
multi_slice_viewer_debug(torch.cat([image1, image2], dim=0),
torch.cat([target1, target2], dim=0),
no_contour=True)
plt.show()
def __init__(self, datadir, resize, argument=True):
self.datadir = datadir
self.pairlist = self.search_pairs()
if argument == False:
self.image_transforms = Compose([
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='springface')
if resize > 0:
self.joint_transforms = ComposeX([
Resize(resize),
])
else:
self.joint_transforms = None
else:
self.image_transforms = Compose([
ColorJitter(0.05, 0.05, 0.05),
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='springface')
if resize > 0:
self.joint_transforms = ComposeX(
[Resize(resize), RandomHorizontallyFlip()])
else:
self.joint_transforms = RandomHorizontallyFlip()
def __init__(self, datadir, resize, crop_size, argument=True):
self.datadir = datadir
self.pairlist = self.search_pairs()
if resize > 0:
f_resize = Resize(resize)
else:
f_resize = None
if crop_size is not None:
f_crop = RandomCrop(crop_size)
else:
f_crop = None
if argument == False:
self.image_transforms = Compose([
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='springhair')
self.joint_transforms = ComposeX([f_resize, f_crop])
else:
self.image_transforms = Compose([
ColorJitter(0.1, 0.1, 0.1),
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='springhair')
self.joint_transforms = ComposeX(
[f_resize, f_crop, RandomHorizontallyFlip()])
def __call__(self, input, target):
## do something to both images
input = Resize((1086, 1351), Image.BILINEAR)(input)
target = Resize((1086, 1351), Image.NEAREST)(target)
#input = Resize((512,1024), Image.BILINEAR)(input)
#target = Resize((512,1024),Image.NEAREST)(target)
if (self.augment):
rotation_degree = 1
shear_degree = 1
input = RandomAffine(rotation_degree,
None,
None,
shear_degree,
resample=Image.BILINEAR,
fillcolor=0)(input)
target = RandomAffine(rotation_degree,
None,
None,
shear_degree,
resample=Image.NEAREST,
fillcolor=255)(target)
w, h = input.size
nratio = random.uniform(0.5, 1.0)
ni = random.randint(0, int(h - nratio * h))
nj = random.randint(0, int(w - nratio * w))
input = input.crop(
(nj, ni, int(nj + nratio * w), int(ni + nratio * h)))
target = target.crop(
(nj, ni, int(nj + nratio * w), int(ni + nratio * h)))
input = Resize((512, 1024), Image.BILINEAR)(input)
target = Resize((512, 1024), Image.NEAREST)(target)
brightness = 0.1
contrast = 0.1
saturation = 0.1
hue = 0.1
input = ColorJitter(brightness, contrast, saturation, hue)(input)
hflip = random.random()
if (hflip < 0.5):
input = input.transpose(Image.FLIP_LEFT_RIGHT)
target = target.transpose(Image.FLIP_LEFT_RIGHT)
else:
input = Resize((512, 1024), Image.BILINEAR)(input)
target = Resize((512, 1024), Image.NEAREST)(target)
input = ToTensor()(input)
if (self.enc):
target = Resize((64, 128), Image.NEAREST)(target)
target = ToLabel()(target)
target = Relabel(255, 27)(target)
return input, target
def __init__(self,
img_size=224,
src_perturbation=0.2,
dst_perturbation=0.05,
brightness=0.3,
contrast=0.5,
saturation=0.1,
hue=0.3):
self.size = img_size
self.src_per = src_perturbation
self.tgt_per = dst_perturbation
self.thre = 50
self.src_const = np.array([
[-0.5, -0.5],
[0.5, -0.5],
[-0.5, 0.5],
[0.5, 0.5],
[0.2, -0.2],
[-0.2, 0.2],
[0.2, 0.2],
[-0.2, -0.2]
])
self.jittering = ColorJitter(brightness, contrast, saturation, hue)
def __init__(self, train_cfg, database):
super().__init__()
self.base_scale = train_cfg['sample_scale_inter']
self.base_rotate = train_cfg['sample_rotate_inter'] / 180 * np.pi
self.database = database
self.transformer = TransformerCV(train_cfg)
self.args = train_cfg['augmentation_args']
self.jitter = ColorJitter(self.args['brightness'], self.args['contrast'], self.args['saturation'],
self.args['hue'])
# img_dir = os.path.join('data', 'SUN2012Images', 'JPEGImages')
#self.background_pths = [os.path.join(img_dir, fn) for fn in os.listdir(img_dir)]
self.name2func = {
'jpeg': lambda img_in: jpeg_compress(img_in, self.args['jpeg_low'], self.args['jpeg_high']),
'blur': lambda img_in: gaussian_blur(img_in, self.args['blur_range']),
'jitter': lambda img_in: np.asarray(self.jitter(Image.fromarray(img_in))),
'noise': lambda img_in: add_noise(img_in),
'none': lambda img_in: img_in,
'sp_gaussian_noise': lambda img_in: additive_gaussian_noise(img_in, self.args['sp_gaussian_range']),
'sp_speckle_noise': lambda img_in: additive_speckle_noise(img_in, self.args['sp_speckle_prob_range']),
'sp_additive_shade': lambda img_in: additive_shade(img_in, self.args['sp_nb_ellipse'], self.
args['sp_transparency_range'],
self.args['sp_kernel_size_range']),
'sp_motion_blur': lambda img_in: motion_blur(img_in, self.args['sp_max_kernel_size']),
'resize_blur': lambda img_in: resize_blur(img_in, self.args['resize_blur_min_ratio'])
}
def __getitem__(self, index):
if self.augment_data == False:
img = cv2.resize(
cv2.imread(self.train_images +
'{}.jpg'.format(self.img_id[index])), (512, 512))
else:
img = Image.open(self.train_images +
'{}.jpg'.format(self.img_id[index])).resize(
(512, 512))
color = ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.6,
hue=0.4)
blur = GaussianBlur((3, 3), sigma=(0.1, 2))
img = color(img)
img = blur(img)
img = np.array(img)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = self.normalize(img)
img = img.transpose([2, 0, 1])
heatmap, offset_x, offset_y, object_size_x, object_size_y = generate_heatmap_offset(
self.img_id[index], self.dictionnary_labels_per_image)
regr = np.zeros((2, 128, 128))
offset = np.zeros((2, 128, 128))
regr[0, :, :] = object_size_x
regr[1, :, :] = object_size_y
offset[0, :, :] = offset_x
offset[1, :, :] = offset_y
return img, self.img_id[index], heatmap, regr, offset
def __init__(self, datadir, resize, argument=True):
self.img_dir = os.path.join(datadir, 'CelebA-HQ-img')
self.mask_dir = os.path.join(datadir, 'CelebAMask-HQ-mask-anno')
if resize < 0:
self.resize = 512
else:
self.resize = resize
if argument == False:
self.image_transforms = Compose([
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='celebamask-hq')
self.joint_transforms = Resize(self.resize)
else:
self.image_transforms = Compose([
ColorJitter(0.05, 0.05, 0.05),
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225])
])
self.mask_transforms = ToIndex(dataset='celebamask-hq')
self.joint_transforms = ComposeX(
[Resize(self.resize),
RandomHorizontallyFlip()])
def __init__(self,
crop_size,
min_scale=-0.2,
max_scale=0.5,
do_flip=False):
# spatial augmentation params
self.crop_size = crop_size
self.min_scale = min_scale
self.max_scale = max_scale
self.spatial_aug_prob = 0.8
self.stretch_prob = 0.8
self.max_stretch = 0.2
# flip augmentation params
self.do_flip = do_flip
self.h_flip_prob = 0.5
self.v_flip_prob = 0.1
# photometric augmentation params
self.photo_aug = ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.3,
hue=0.3 / 3.14)
self.asymmetric_color_aug_prob = 0.2
self.eraser_aug_prob = 0.5
def __init__(self, phase, kwargs):
self.mode = Mode[kwargs['mode']]
self.image_size = kwargs['image_size']
self.hidden_size = kwargs['hidden_size']
self.debug_use_dataset = kwargs['debug_use_dataset']
self.debug_one_sentence = kwargs['debug_one_sentence']
self.__use_densenet = kwargs['__use_densenet']
self.sent_tokenizer = PunktSentenceTokenizer()
self.word_tokenizer = TweetTokenizer()
if phase == Phase.train:
jitter = [ColorJitter(brightness=0.5, contrast=0.5)]
else:
jitter = []
if self.__use_densenet:
self.transform = Compose((
[Lambda(lambda img: img.convert('RGB'))] +
[Resize((256, 256))] +
jitter +
[ToTensor()] +
[Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
))
else:
self.transform = Compose((
[Resize((256, 256))] +
jitter +
[ToTensor()]
))
def __init__(self,
image_folder,
max_images=False,
image_size=(512, 512),
add_random_masks=False):
super(ImageInpaintingData, self).__init__()
if isinstance(image_folder, str):
self.images = glob.glob(os.path.join(image_folder, "raw/*"))
else:
self.images = list(
chain.from_iterable([
glob.glob(os.path.join(i, "raw/*")) for i in image_folder
]))
assert len(self.images) > 0
if max_images:
self.images = random.choices(self.images, k=max_images)
print(f"Find {len(self.images)} images.")
self.img_size = image_size
self.transformer = Compose([
RandomGrayscale(p=0.4),
ColorJitter(brightness=0.2, contrast=0.2, saturation=0, hue=0),
ToTensor(),
# Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
self.add_random_masks = add_random_masks
self.random_mask = RandomMask(image_size[0])
def test_color_jitter(brightness, contrast, saturation, hue):
np.random.seed(0)
img = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
pil_image = Image.fromarray(img)
transform = A.Compose([
A.ColorJitter(
brightness=[brightness, brightness],
contrast=[contrast, contrast],
saturation=[saturation, saturation],
hue=[hue, hue],
p=1,
)
])
pil_transform = ColorJitter(
brightness=[brightness, brightness],
contrast=[contrast, contrast],
saturation=[saturation, saturation],
hue=[hue, hue],
)
res1 = transform(image=img)["image"]
res2 = np.array(pil_transform(pil_image))
_max = np.abs(res1.astype(np.int16) - res2.astype(np.int16)).max()
assert _max <= 2, "Max: {}".format(_max)
def maybe_darken_a_lot(image):
if randint(0, 100) <= 30:
brightness = uniform(.5, .8)
saturation = uniform(1., 1.5)
return ColorJitter(brightness=(brightness, brightness),
saturation=(saturation, saturation))(image)
return image
def transform_image_only(image):
if random.random() > 0.5:
# random change the colors of the picture
image = RandomChoice(
[Grayscale(num_output_channels=3),
ColorJitter(1, 1, 1)])(image)
return image
def get_train_transforms(img_size: int) -> Compose:
"""Returns data transformations/augmentations for train dataset.
Args:
img_size: The resolution of the input image (img_size x img_size)
"""
return Compose([
RandomApply([
ColorJitter(brightness=0.3, contrast=0.01, saturation=0.01, hue=0),
RandomAffine(0.1,
translate=(0.04, 0.04),
scale=(0.04, 0.04),
shear=0.01,
resample=2),
#Grayscale(num_output_channels=3),
#RandomCrop(30),
RandomPerspective(0.1)
]),
Resize([img_size, img_size], interpolation=3),
ToTensor(),
#RandomApply([
#RandomErasing(p=0.2, scale=(0.02, 0.33), ratio=(0.3, 3.3))
#]),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def cifar10(root: str,
batch_size: int = 32,
workers: int = 6,
augment: bool = True,
splits: Union[str, Tuple[str]] = ('train', 'val')) -> LoaderTypes:
"""Wrapper for loading the `CIFAR10` dataset.
Args:
root: The root directory where the dataset is stored. Usually ~/.torch/datasets.
batch_size: The batch size.
workers: The number of CPUs to use for when loading the data from disk.
augment: Whether to use data augmentation when training.
splits: Which splits of the data to return. Possible values are `train`, `val` and `test`.
Returns:
A list data loaders of the chosen splits.
"""
normalize = Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
val_transform = Compose([ToTensor(), normalize])
if augment:
transform = Compose(
[RandomCrop(32, padding=4),
RandomHorizontalFlip(),
ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
RandomRotation(degrees=5),
ToTensor(),
normalize])
else:
transform = val_transform
loader_list = list()
if 'train' in splits:
train_val_set = CIFAR10(root, train=True, transform=transform, download=True)
train_loader = DataLoader(train_val_set, batch_size=batch_size, shuffle=True, num_workers=workers,
pin_memory=True)
loader_list.append(train_loader)
if 'test' in splits or 'val' in splits:
val_test_set = CIFAR10(root, train=False, transform=val_transform, download=True)
val_set, test_set = torch.utils.data.random_split(val_test_set, [5000, 5000])
if 'val' in splits:
val_set = Memory(val_set, img_size=32, channels=3)
for _ in val_set:
pass
val_set.set_use_cache(True)
val_set.pin_memory()
loader_list.append(val_set)
if 'test' in splits:
test_set = Memory(test_set, img_size=32, channels=3)
for _ in test_set:
pass
test_set.set_use_cache(True)
test_set.pin_memory()
loader_list.append(test_set)
if len(loader_list) == 1:
return loader_list[0]
return loader_list
