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