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 get_transform(new_size=None, augment_images=False):
"""
obtain the image transforms required for the input data
aurora images are already flipped so perform random rotation instead
:param new_size: size of the resized images
:param augment_images: Whether to randomly mirror input images during training
:return: image_transform => transform object from TorchVision
"""
from torchvision.transforms import ToTensor, Normalize, Compose, Resize, \
RandomHorizontalFlip, RandomAffine
print('get_transform(): new_size={}, augment_images={}'.format(
new_size, augment_images))
if not augment_images:
if new_size is not None:
print('get_transform(): A')
image_transform = Compose([
Resize(new_size),
ToTensor(),
Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
else:
print('get_transform(): B')
image_transform = Compose([
ToTensor(),
Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
else:
if new_size is not None:
print('get_transform(): C')
image_transform = Compose([
RandomAffine(10),
Resize(new_size),
ToTensor(),
Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
else:
print('get_transform(): D')
image_transform = Compose([
RandomAffine(10),
ToTensor(),
Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
return image_transform
def __init__(self):
self.path_images = folder_images
file_list = glob(self.path_images + "*")
file_list.sort()
labels = get_labels_train() ## labels are sorted already
self.test_data = False
self.data = []
for file, label in zip(file_list, labels):
self.data.append([file, label])
self.augment = Compose([
# RandomVerticalFlip(),
RandomHorizontalFlip(),
RandomAffine((0, 360),
(0.01, ) * 2), # rotation, -4 to 4 translation
CenterCrop(207),
# Resize((150,)*2),
# Resize((100,)*2),
Resize((64, ) * 2),
ToTensor(),
Normalize(0.5, 0.5),
])
self.augment_test = Compose([
# RandomVerticalFlip(),
# RandomHorizontalFlip(),
# RandomAffine((0,360), (0.01,)*2), # rotation, -4 to 4 translation
CenterCrop(207),
# Resize((150,)*2),
# Resize((100,)*2),
Resize((64, ) * 2),
ToTensor(),
Normalize(0.5, 0.5),
])
def load_train_and_evaluation_img(batch_size=64):
random_split_train_test(os.path.dirname(__file__) + "/data/train.csv",
train_rate=0.75)
train_transform = transforms.Compose([
transforms.ToPILImage(),
RandomAffine(degrees=20, translate=(0.1, 0.2), scale=(.9, 1.1)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, ), std=(0.5, ))
])
test_transform = transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, ), std=(0.5, ))
])
train_dataset = MNIST_data(
os.path.dirname(__file__) + "/data/train_p1.csv", IMG_PIXELS,
IMG_HEIGHT, IMG_WIDTH, train_transform)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_dataset = MNIST_data(
os.path.dirname(__file__) + "/data/train_p2.csv", IMG_PIXELS,
IMG_HEIGHT, IMG_WIDTH, test_transform)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
return train_loader, test_loader
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 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(),
Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_transform = Compose([
Resize((img_size, img_size)),
ToTensor(),
Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
train_loader = DataLoader(HairStyleDataset(train_files, train_transform),
batch_size=BATCH_SIZE,
shuffle=True,
pin_memory=True,
num_workers=4)
val_loader = DataLoader(HairStyleDataset(val_files, val_transform),
batch_size=BATCH_SIZE,
shuffle=True,
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 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_image_transform(self) -> ModelTransformsPerExecutionMode:
if self.use_combined_model:
return ModelTransformsPerExecutionMode(
train=ImageTransformationPipeline(
transforms=[RandomAffine(degrees=30, translate=(0.1, 0.1), shear=15),
ColorJitter(brightness=0.2)]))
else:
return ModelTransformsPerExecutionMode()
def __init__(self, img_loc, parent_mask_loc, mask_type, total_data_number,
reuse_parent_mask):
self.img_loc = img_loc
self.parent_mask_loc = parent_mask_loc
self.total_data_number = total_data_number
self.reuse_parent_mask = reuse_parent_mask
self.augmentation = RandomAffine(degrees=180, scale=[0.7, 1.2])
self.mask_generator = self._mask_fn(mask_type)
def get_augmentation(augmentation, dataset, data_shape):
c, h, w = data_shape
if augmentation is None:
pil_transforms = []
elif augmentation == 'horizontal_flip':
pil_transforms = [RandomHorizontalFlip(p=0.5)]
elif augmentation == 'neta':
assert h==w
pil_transforms = [Pad(int(math.ceil(h * 0.04)), padding_mode='edge'),
RandomAffine(degrees=0, translate=(0.04, 0.04)),
CenterCrop(h)]
elif augmentation == 'eta':
assert h==w
pil_transforms = [RandomHorizontalFlip(),
Pad(int(math.ceil(h * 0.04)), padding_mode='edge'),
RandomAffine(degrees=0, translate=(0.04, 0.04)),
CenterCrop(h)]
return pil_transforms
def get_rotations():
transformations = {}
for name, angle in ROTATIONS.items():
trans = transforms.Compose([
transforms.ToPILImage(),
RandomAffine(degrees=(angle, angle)),
transforms.ToTensor()
])
transformations[name] = BatchedTransformation(trans)
return transformations
def get_image_transform(self) -> ModelTransformsPerExecutionMode:
"""
Get transforms to perform on image samples for each model execution mode.
"""
if self.imaging_feature_type in [ImagingFeatureType.Image, ImagingFeatureType.ImageAndSegmentation]:
return ModelTransformsPerExecutionMode(
train=ImageTransformationPipeline(
transforms=[RandomAffine(10), ColorJitter(0.2)],
use_different_transformation_per_channel=True))
return ModelTransformsPerExecutionMode()
def __call__(self, image):
rand_ = random.uniform(0, 1)
if rand_ < self.prob:
RandAffine_ = RandomAffine(degrees=self.degrees,
translate=(self.translate,
self.translate),
scale=self.scale,
shear=self.shear)
image = RandAffine_(image)
return image
def get_scales():
transformations = {}
for name, scale in SCALES.items():
trans = transforms.Compose([
transforms.ToPILImage(),
RandomAffine(degrees=0, scale=(scale, scale)),
transforms.ToTensor()
])
transformations[name] = BatchedTransformation(trans)
return transformations
def __init__(self, data_dir, batch_size, num_workers):
super().__init__()
transforms = Compose([
ToTensor(),
RandomHorizontalFlip(),
RandomAffine(degrees=5, translate=(0.05, 0.05), scale=(1., 1.5))
])
self.dataset = UTKFaceDataset(data_dir, transforms)
self.batch_size = batch_size
self.num_workers = num_workers
def create_dataloaders(Dataset,
base_size=32,
crop_size=28,
batch_size=4,
supervised_proportion=1.0):
train_tfms = Compose([
RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.7, 1.3)),
ToTensor(),
Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.247, 0.243, 0.261))
])
val_tfms = Compose([
ToTensor(),
Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.247, 0.243, 0.261))
])
train_dataset = Dataset('data',
train=True,
transform=train_tfms,
download=True)
val_dataset = Dataset('data',
train=False,
transform=val_tfms,
download=True)
# For the training dataset, we need to split the indices into
# labelled and unlabelled indices, then use TwoStreamBatchSampler
if supervised_proportion < 1.0:
np.random.seed(0)
labeled_idxs = np.hstack(
list(
yield_proportion_of_labelled_indices(
train_dataset, np.arange(len(train_dataset)),
supervised_proportion)))
unlabeled_idxs = np.setdiff1d(np.arange(len(train_dataset)),
labeled_idxs)
# Mutate the training dataset in-place. Not ideal, but this is
# at least consistent with two-moons
targets = np.array(train_dataset.train_labels)
targets[unlabeled_idxs] = -1
train_dataset.train_labels = list(targets)
return (DataLoader(train_dataset,
batch_sampler=TwoStreamBatchSampler(
primary_indices=labeled_idxs,
secondary_indices=unlabeled_idxs,
batch_size=batch_size,
secondary_batch_size=int(
batch_size * (1 - supervised_proportion))),
pin_memory=True),
DataLoader(val_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True))
def __init__(self, degrees, interpolation):
assert isinstance(degrees, float)
if interpolation == "NEAREST":
interpolation = InterpolationMode.NEAREST
elif interpolation == "BILINEAR":
interpolation = InterpolationMode.BILINEAR
else:
raise TypeError(f"Transform interpolation {interpolation} not supported.")
self.rotate = RandomAffine(degrees,
interpolation=interpolation)
self.pad = None
def SR4x_transform(subset):
try:
ts = Compose([
RandomCrop(
(subset,
subset)), #, pad_if_needed=False, padding_mode='reflect'),
RandomHorizontalFlip(),
RandomVerticalFlip(),
ColorJitter(hue=.05, saturation=.05),
RandomAffine(30),
RandomGrayscale(),
])
except:
ts = Compose([
Resize(subset),
RandomHorizontalFlip(),
RandomVerticalFlip(),
ColorJitter(hue=.05, saturation=.05),
RandomAffine(30),
RandomGrayscale(),
])
return ts
def random_affine_transformation(image):
"""Applies a random affine transformation to the image."""
transform = Compose([
RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.8, 1.3), shear=10),
Resize(32),
ToTensor(),
Normalize(mean=(0.5,), std=(0.5,))
])
transform_deepfashion = Compose([
RandomAffine(degrees=15, translate=(0.1, 0.1), scale=(0.8, 1.3), shear=10),
Resize(32),
ToTensor(),
Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
img = None
if image.mode == 'L':
img = transform(image)
if image.mode == 'RGB':
img = transform_deepfashion(image)
img = img.to(DEVICE)
return img
def load_train_img(batch_size=64):
train_img_dir = os.path.dirname(__file__) + "/data/train.csv"
train_transform = transforms.Compose([
transforms.ToPILImage(),
RandomAffine(degrees=20, translate=(0.1, 0.2), scale=(.9, 1.1)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, ), std=(0.5, ))
])
train_dataset = MNIST_data(train_img_dir, IMG_PIXELS, IMG_HEIGHT,
IMG_WIDTH, train_transform)
return torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
worker_init_fn=seed_init_fn)
def get_transform(
target_size=(512, 512),
transform_list='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.6, 1.0) if is_train else (0.8, 1.0)
transform.append(RandomResizedCrop(target_size, scale=scale))
# elif transform_name == 'resize':
# transform.append(Resize(target_size))
elif transform_name == 'keep_aspect':
transform.append(KeepAsepctResize(target_size))
elif transform_name == 'Affine':
augments.append(
RandomAffine(degrees=(-180, 180),
scale=(0.8889, 1.0),
shear=(-36, 36)))
elif transform_name == 'centor_crop':
augments.append(CenterCrop(target_size))
elif transform_name == 'horizontal_flip':
augments.append(RandomHorizontalFlip())
elif transform_name == 'vertical_flip':
augments.append(RandomVerticalFlip())
elif transform == 'random_grayscale':
p = 0.5 if is_train else 0.25
transform.append(RandomGrayscale(p))
elif transform_name == 'random_rotate':
augments.append(RandomRotation(180))
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))
transform.append(ToTensor())
transform.append(
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
return Compose(transform)
def __init__(self):
image_size = 224
train_transform = Compose([
Resize(image_size, BICUBIC),
RandomAffine(degrees=2,
translate=(0.02, 0.02),
scale=(0.98, 1.02),
shear=2,
fillcolor=(124, 117, 104)),
RandomHorizontalFlip(),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_transform = Compose([
Resize(image_size, BICUBIC),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
self.summaryWriter = SummaryWriter("./logs")
# Dataloader装载训练数据集，batch_size每轮100个数据，shuffle并打乱顺序
self.train_dataset = datasets.CIFAR10("../data/CIFAR10/",
True,
transform=train_transform,
download=True)
# Dataloader装载测试数据集，batch_size每轮100个数据，shuffle并打乱顺序
self.train_dataload = DataLoader(self.train_dataset,
batch_size=100,
shuffle=True)
self.test_dataset = datasets.CIFAR10("../data/CIFAR10/",
False,
transform=test_transform,
download=True)
self.test_dataload = DataLoader(self.test_dataset,
batch_size=100,
shuffle=True)
# 创建网络
self.net = NetV2()
# 装载之前训练状态
# self.net.load_state_dict(torch.load("./checkpoint/27.t"))
# 将数据移动至GPU运算
# self.net.to(DEVICE)
# 创建优化器，将网络中net.parameters()参数放入优化器
self.opt = optim.Adam(self.net.parameters())
self.loss_fn = nn.CrossEntropyLoss()
def __init__(self,
model=None,
defense_name=None,
dataset=None,
re_training=True,
training_parameters=None,
device=None,
**kwargs):
super(EITDefense, self).__init__(model=model,
defense_name=defense_name)
self.model = model
self.defense_name = defense_name
self.device = device
self.Dataset = dataset.upper()
assert self.Dataset in ['MNIST', 'CIFAR10'
], "The data set must be MNIST or CIFAR10"
# make sure to parse the parameters for the defense
assert self._parsing_parameters(**kwargs)
if re_training:
# get the training_parameters, the same as the settings of RawModels
self.num_epochs = training_parameters['num_epochs']
self.batch_size = training_parameters['batch_size']
# prepare the optimizers
if self.Dataset == 'MNIST':
self.optimizer = optim.SGD(
self.model.parameters(),
lr=training_parameters['learning_rate'],
momentum=training_parameters['momentum'],
weight_decay=training_parameters['decay'],
nesterov=True)
else:
self.optimizer = optim.Adam(self.model.parameters(),
lr=training_parameters['lr'])
self.color_mode = 'RGB' if self.Dataset == 'CIFAR10' else 'L'
if self.Dataset == 'CIFAR10':
self.transform = Compose([
RandomAffine(degrees=0, translate=(0.1, 0.1)),
RandomHorizontalFlip(),
ToTensor()
])
else:
self.transform = None
def __getitem__(self, idx):
if not self.load_into_memory:
img_tensor, mask_tensor = self._load_data(self.data_paths[idx], self.mask_paths[idx])
else:
mask_tensor = self.masks[idx]
img_tensor = self.data[idx]
if self.data_augmentation:
img = self.toPIL(img_tensor)
transform = RandomAffine.get_params(self.aug_rotation, self.aug_translation, self.aug_scale, None, img.size)
img_tensor = self.totensor(affine(img, *transform, resample=False, fillcolor=0))
mask_tensor = self.totensor(affine(self.toPIL(mask_tensor), *transform, resample=False, fillcolor=0))
if self.randomized_background:
img_tensor = img_tensor * mask_tensor
if self.randomized_background:
if not self.load_into_memory:
bg_path = self.background_paths[np.random.randint(0, len(self.backgrounds))]
bg_img = Image.open(bg_path)
#bg_img = bg_img.resize((self.image_width, self.image_height))
bg_tensor = self.totensor(bg_img)
else:
bg_tensor = self.backgrounds[np.random.randint(0, len(self.backgrounds))]
bg_tensor = self.totensor(self.random_crop(self.toPIL(bg_tensor)))
inverse_mask = torch.ones(mask_tensor.shape, dtype=torch.float) - mask_tensor
masked_background = bg_tensor * inverse_mask
img_tensor = img_tensor + masked_background
if self.visualize_data:
img_vis = img_tensor - torch.min(img_tensor)
img_vis = img_vis / torch.max(img_vis)
img = self.toPIL(img_vis)
img.save(os.path.join(self.output_dir, 'train', 'example_train_image_{}.jpg'.format(idx)))
return img_tensor, self.labels[idx]
def test_augment():
from torchvision.transforms import RandomApply, ColorJitter, RandomAffine, ToTensor, Normalize
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
fill_color = (255, 255, 255) # (209, 200, 193)
min_scale, max_scale = 2 / 3, 2
imgW = 1024
imgH = 64
transform_train = transforms.Compose([
# RandomApply([cnx_aug_thin_characters()], p=0.2),
# RandomApply([cnx_aug_bold_characters()], p=0.4),
# cnd_aug_randomResizePadding(imgH, imgW, min_scale, max_scale, fill=fill_color),
cnd_aug_resizePadding(imgW, imgH, fill=fill_color),
RandomApply([cnd_aug_add_line()], p=0.5),
RandomApply([cnx_aug_blur()], p=0.3),
ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
RandomApply(
[
RandomAffine(
shear=(-20, 20),
translate=(0.0, 0.1),
degrees=0,
# degrees=2,
# scale=(0.8, 1),
fillcolor=fill_color)
],
p=0.5)
# ,ToTensor()
# ,Normalize(mean, std)
])
path = "/data/train_data_29k_29Feb_update30Mar/cinnamon_data/cinamon_test_115/"
# path='/data/SDV/cropped_img'
list_img = list_files1(path, "jpg")
print(list_img)
for l in list_img:
path_img = os.path.join(path, l)
img = cv2.imread(path_img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
img_rs = transform_train(img)
numpy_image = np.array(img_rs)
# numpy_image = augment_random_rotate(img, -5, 5)
cv_img = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
cv2.imshow("result", cv_img)
cv2.waitKey(0)
def augmentation_transformer_factory():
degrees, translation = augmentation_config
# For some reasons fillcolor does not work therefore we substract the min value before and add it after
r = RandomAffine(degrees=degrees,
shear=degrees,
translate=translation,
resample=PIL.Image.BILINEAR)
def transformer(data):
target = data.detach().cpu().numpy().copy() - image_min_value
for i in range(0, data.shape[0]):
image = target[i, :].reshape((image_wh, image_wh))
transformed = r(PIL.Image.fromarray(image))
transformed = np.array(transformed)
target[i, :] = transformed.reshape((1, image_wh * image_wh))
return torch.from_numpy(target + image_min_value).to(data.device)
return transformer
class SyncRandomAffine():
def __init__(self, degrees, translate=None, scale=None, shear=None):
self.trans = RandomAffine(degrees, translate, scale, shear)
def __call__(self, img, mask):
ret = self.trans.get_params(self.trans.degrees, self.trans.translate,
self.trans.scale, self.trans.shear,
img.size)
img_prime = F.affine(img,
*ret,
interpolation=self.trans.resample,
fill=self.trans.fillcolor)
mask_prime = F.affine(mask,
*ret,
interpolation=self.trans.resample,
fill=self.trans.fillcolor)
return img_prime, mask_prime
def test_torchvision_on_various_input(
use_different_transformation_per_channel: bool,
) -> None:
"""
This tests that we can run transformation pipeline with out of the box torchvision transforms on various types
of input: PIL image, 3D tensor, 4D tensors. Tests that use_different_transformation_per_channel has the correct
behavior.
"""
image_as_pil, image_2d_as_CHW_tensor, image_2d_as_ZCHW_tensor, scan_4d_as_tensor = create_test_images()
transform = ImageTransformationPipeline(
[
CenterCrop(crop_size),
RandomErasing(),
RandomAffine(degrees=(10, 12), shear=15, translate=(0.1, 0.3)),
],
use_different_transformation_per_channel,
)
# Test PIL image input
transformed = transform(image_as_pil)
assert isinstance(transformed, torch.Tensor)
assert transformed.shape == torch.Size([1, crop_size, crop_size])
# Test image as [C, H. W] tensor
transformed = transform(image_2d_as_CHW_tensor.clone())
assert isinstance(transformed, torch.Tensor)
assert transformed.shape == torch.Size([1, crop_size, crop_size])
# Test image as [1, 1, H, W]
transformed = transform(image_2d_as_ZCHW_tensor)
assert isinstance(transformed, torch.Tensor)
assert transformed.shape == torch.Size([1, 1, crop_size, crop_size])
# Test with a fake 4D scan [C, Z, H, W] -> [25, 34, 32, 32]
transformed = transform(scan_4d_as_tensor)
assert isinstance(transformed, torch.Tensor)
assert transformed.shape == torch.Size([5, 4, crop_size, crop_size])
# Same transformation should be applied to all slices and channels.
assert (
torch.isclose(transformed[0, 0], transformed[1, 1]).all()
!= use_different_transformation_per_channel
)
def random_rotation(im, hm, degrees, scale=(0.8, 1.2)):
topil = ToPILImage()
totensor = ToTensor()
rot = RandomAffine(degrees=degrees, scale=scale)
seed = np.random.randint(0, 2 ** 32)
im_pil = topil(im)
random.seed(seed)
im_r = rot(im_pil)
im = totensor(im_r)
hm_new = torch.zeros(hm.size())
for i in range(hm.size(0)):
hm_pil = topil(hm[i, :, :].unsqueeze(0))
random.seed(seed)
hm_r = rot((hm_pil))
hm_t = totensor(hm_r).squeeze()
hm_new[i, :, :] = hm_t
return im, hm_new