Source code for tensornet.data.processing
import numpy as np
import torch
import albumentations as A
from albumentations.pytorch import ToTensor
[docs]class Transformations:
"""Wrapper class to pass on albumentaions transforms into PyTorch."""
def __init__(
self, resize=(0, 0), padding=(0, 0), crop=(0, 0), horizontal_flip_prob=0.0,
vertical_flip_prob=0.0, gaussian_blur_prob=0.0, rotate_degree=0.0,
cutout_prob=0.0, cutout_dim=(8, 8), hue_saturation_prob=0.0, contrast_prob=0.0,
mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True, train=True
):
"""Create data transformation pipeline.
Args:
resize (:obj:`tuple`, optional): Resize the input to the given height and
width. (default: (0, 0))
padding (:obj:`tuple`, optional): Pad the image if the image size is less
than the specified dimensions (height, width). (default= (0, 0))
crop (:obj:`tuple`, optional): Randomly crop the image with the specified
dimensions (height, width). (default: (0, 0))
horizontal_flip_prob (:obj:`float`, optional): Probability of an image
being horizontally flipped. (default: 0)
vertical_flip_prob (:obj:`float`, optional): Probability of an image
being vertically flipped. (default: 0)
rotate_degree (:obj:`float`, optional): Angle of rotation for image
augmentation. (default: 0)
cutout_prob (:obj:`float`, optional): Probability that cutout will be
performed. (default: 0)
cutout_dim (:obj:`tuple`, optional): Dimensions of the cutout box (height, width).
(default: (8, 8))
hue_saturation_prob (:obj:`float`, optional): Probability of randomly changing hue,
saturation and value of the input image. (default: 0)
contrast_prob (:obj:`float`, optional): Randomly changing contrast of the input image.
(default: 0)
mean (:obj:`float` or :obj:`tuple`, optional): Dataset mean. (default: 0.5 for each channel)
std (:obj:`float` or :obj:`tuple`, optional): Dataset standard deviation.
(default: 0.5 for each channel)
"""
transforms_list = []
if sum(resize) > 0:
transforms_list += [A.Resize(
height=resize[0], width=resize[1], always_apply=True
)]
if train:
if sum(padding) > 0:
transforms_list += [A.PadIfNeeded(
min_height=padding[0], min_width=padding[1], always_apply=True
)]
if sum(crop) > 0:
transforms_list += [A.RandomCrop(crop[0], crop[1], always_apply=True)]
if horizontal_flip_prob > 0: # Horizontal Flip
transforms_list += [A.HorizontalFlip(p=horizontal_flip_prob)]
if vertical_flip_prob > 0: # Vertical Flip
transforms_list += [A.VerticalFlip(p=vertical_flip_prob)]
if gaussian_blur_prob > 0: # Patch Gaussian Augmentation
transforms_list += [A.GaussianBlur(p=gaussian_blur_prob)]
if rotate_degree > 0: # Rotate image
transforms_list += [A.Rotate(limit=rotate_degree)]
if cutout_prob > 0: # CutOut
if isinstance(mean, float):
fill_value = mean * 255.0
else:
fill_value = tuple([x * 255.0 for x in mean])
transforms_list += [A.CoarseDropout(
p=cutout_prob, max_holes=1, fill_value=fill_value,
max_height=cutout_dim[0], max_width=cutout_dim[1]
)]
if hue_saturation_prob > 0: # Hue Saturation
transforms_list += [A.HueSaturationValue(p=hue_saturation_prob)]
if contrast_prob > 0: # Random Contrast
transforms_list += [A.RandomContrast(p=contrast_prob)]
if normalize:
# normalize the data with mean and standard deviation to keep values in range [-1, 1]
# since there are 3 channels for each image,
# we have to specify mean and std for each channel
transforms_list += [
A.Normalize(mean=mean, std=std, always_apply=True),
]
# convert the data to torch.FloatTensor
transforms_list += [
ToTensor()
]
self.transform = A.Compose(transforms_list)
[docs] def __call__(self, image):
"""Process and image through the data transformation pipeline.
Args:
image: Image to process.
Returns:
(*torch.Tensor*): Transformed image.
"""
if not isinstance(image, np.ndarray):
image = np.array(image)
if len(image.shape) == 2:
image = np.expand_dims(image, axis=-1)
image = self.transform(image=image)['image']
return image
[docs]def data_loader(data, shuffle=True, batch_size=1, num_workers=1, cuda=False):
"""Create data loader
Args:
data (torchvision.datasets): Downloaded dataset.
shuffle (:obj:`bool`, optional): If True, shuffle the dataset.
(default: True)
batch_size (:obj:`int`, optional): Number of images to considered
in each batch. (default: 1)
num_workers (:obj:`int`, optional): How many subprocesses to use
for data loading. (default: 1)
cuda (:obj:`bool`, optional): True is GPU is available. (default: False)
Returns:
(*torch.utils.data.DataLoader*): DataLoader instance.
"""
loader_args = {
'shuffle': shuffle,
'batch_size': batch_size
}
# If GPU exists
if cuda:
loader_args['num_workers'] = num_workers
loader_args['pin_memory'] = True
return torch.utils.data.DataLoader(data, **loader_args)
[docs]class InfiniteDataLoader:
"""Create infinite loop in a data loader.
Args:
data_loader (torch.utils.data.DataLoader): DataLoader object.
auto_reset (:obj:`bool`, optional): Create an infinite loop data loader.
(default: True)
"""
def __init__(self, data_loader, auto_reset=True):
self.data_loader = data_loader
self.auto_reset = auto_reset
self._iterator = iter(data_loader)
def __next__(self):
# Get a new set of inputs and labels
try:
data, target = next(self._iterator)
except StopIteration:
if not self.auto_reset:
raise
self._iterator = iter(self.data_loader)
data, target = next(self._iterator)
return data, target