""" Online Augmentations May 23rd 2023 21:30 ref: CutOut, Mixup, CutMix based on https://blog.csdn.net/cp1314971/article/details/106612060 """ import cv2 import torch import numpy as np import torch.nn.functional as F from scipy.special import perm from torchvision.transforms import Resize from torchvision.transforms import ToPILImage, ToTensor from utils.visual_usage import patchify, unpatchify from utils.fmix import sample_mask, FMixBase # Fmix # generate random bounding box def rand_bbox(size, lam): W = size[2] H = size[3] cut_rat = np.sqrt(1. - lam) cut_w = np.int64(W * cut_rat) cut_h = np.int64(H * cut_rat) # uniform cx = np.random.randint(W) cy = np.random.randint(H) bbx1 = np.clip(cx - cut_w // 2, 0, W) bby1 = np.clip(cy - cut_h // 2, 0, H) bbx2 = np.clip(cx + cut_w // 2, 0, W) bby2 = np.clip(cy + cut_h // 2, 0, H) return bbx1, bby1, bbx2, bby2 def saliency_bbox(img, lam): size = img.size() W = size[1] H = size[2] cut_rat = np.sqrt(1. - lam) cut_w = np.int(W * cut_rat) cut_h = np.int(H * cut_rat) # initialize OpenCV's static fine grained saliency detector and # compute the saliency map temp_img = img.cpu().numpy().transpose(1, 2, 0) saliency = cv2.saliency.StaticSaliencyFineGrained_create() (success, saliencyMap) = saliency.computeSaliency(temp_img) saliencyMap = (saliencyMap * 255).astype("uint8") maximum_indices = np.unravel_index(np.argmax(saliencyMap, axis=None), saliencyMap.shape) x = maximum_indices[0] y = maximum_indices[1] bbx1 = np.clip(x - cut_w // 2, 0, W) bby1 = np.clip(y - cut_h // 2, 0, H) bbx2 = np.clip(x + cut_w // 2, 0, W) bby2 = np.clip(y + cut_h // 2, 0, H) return bbx1, bby1, bbx2, bby2 # augmentation methods class Cutout(object): def __init__(self, alpha=2, shuffle_p=1.0, class_num=2, batch_size=4, device='cpu'): """ Cutout augmentation arXiv:1708.04552 :param alpha: alpha :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param device: CUDA or CPU """ self.alpha = alpha self.class_num = class_num self.batch_size = batch_size self.p = shuffle_p self.device = torch.device(device) def __call__(self, inputs, labels, act=True): labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) # duplicate inputs for ori inputs cutout_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act): # trigger the augmentation operation lam_list.append(-1) continue lam = np.random.beta(self.alpha, self.alpha) bbx1, bby1, bbx2, bby2 = rand_bbox(ori_inputs.size(), lam) # get random bbox cutout_inputs[i, :, bbx1:bbx2, bby1:bby2] = 0 # update the ratio of (area of ori_image on new masked image) for soft-label lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (ori_inputs.size()[2] * ori_inputs.size()[3])) lam_list.append(lam) long_label = labels.argmax(dim=1) # NOTICE cutout use long label and ori_crossentropy instead of soft-label and soft-label_crossentropy return cutout_inputs, long_label, long_label class CutMix(object): def __init__(self, alpha=2, shuffle_p=1.0, class_num=2, batch_size=4, device='cpu'): """ CutMix augmentation arXiv:1905.04899 :param alpha: alpha :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param device: CUDA or CPU """ self.alpha = alpha self.class_num = class_num self.batch_size = batch_size # calibrate the trigger chance of p, new ratio is the change of operation occur in each batch self.p = shuffle_p * (perm(self.batch_size, self.batch_size) / (perm(self.batch_size, self.batch_size) - perm(self.batch_size - 1, self.batch_size - 1))) self.device = torch.device(device) def __call__(self, inputs, labels, act=True): labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) # duplicate inputs for ori inputs cutmix_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio indices = torch.randperm(self.batch_size, device=self.device) # shuffle indices shuffled_inputs = inputs[indices].to(self.device) shuffled_labels = labels[indices].to(self.device) for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act): # trigger the augmentation operation lam_list.append(-1) continue lam = np.random.beta(self.alpha, self.alpha) bbx1, bby1, bbx2, bby2 = rand_bbox(ori_inputs.size(), lam) # get random bbox cutmix_inputs[i, :, bbx1:bbx2, bby1:bby2] = \ shuffled_inputs[i, :, bbx1:bbx2, bby1:bby2] # update the ratio of (area of ori_image on new image) for soft-label lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (ori_inputs.size()[2] * ori_inputs.size()[3])) lam_list.append(lam) labels[i] = labels[i] * lam + shuffled_labels[i] * (1 - lam) long_label = labels.argmax(dim=1) return cutmix_inputs, labels, long_label class Mixup(object): def __init__(self, alpha=2, shuffle_p=1.0, class_num=2, batch_size=4, device='cpu'): """ Mixup augmentation arXiv:1710.09412 :param alpha: alpha :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param device: CUDA or CPU """ self.alpha = alpha self.class_num = class_num self.batch_size = batch_size # calibrate the trigger chance of p, new ratio is the change of operation occur in each batch self.p = shuffle_p * (perm(self.batch_size, self.batch_size) / (perm(self.batch_size, self.batch_size) - perm(self.batch_size - 1, self.batch_size - 1))) self.device = torch.device(device) def __call__(self, inputs, labels, act=True): labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) # duplicate inputs for ori inputs mixup_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio indices = torch.randperm(self.batch_size, device=self.device) # shuffle indices shuffled_inputs = inputs[indices].to(self.device) shuffled_labels = labels[indices].to(self.device) for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act): # trigger the augmentation operation lam_list.append(-1) continue lam = np.random.beta(self.alpha, self.alpha) lam_list.append(lam) mixup_inputs[i] = ori_inputs[i] * lam + shuffled_inputs[i] * (1 - lam) labels[i] = labels[i] * lam + shuffled_labels[i] * (1 - lam) long_label = labels.argmax(dim=1) return mixup_inputs, labels, long_label class SaliencyMix(object): def __init__(self, alpha=1, shuffle_p=1.0, class_num=2, batch_size=4, device='cpu'): """ SaliencyMix augmentation arXiv:2006.01791 :param alpha: alpha :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param device: CUDA or CPU """ # ori batch_size=128 self.alpha = alpha self.class_num = class_num self.batch_size = batch_size # calibrate the trigger chance of p, new ratio is the change of operation occur in each batch self.p = shuffle_p self.device = torch.device(device) def __call__(self, inputs, labels, act=True): labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) # duplicate inputs for ori inputs saliencymix_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio indices = torch.randperm(self.batch_size, device=self.device) # shuffle indices shuffled_inputs = inputs[indices].to(self.device) shuffled_labels = labels[indices].to(self.device) for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act) or self.alpha <= 0: # trigger the augmentation operation lam_list.append(-1) continue lam = np.random.beta(self.alpha, self.alpha) bbx1, bby1, bbx2, bby2 = saliency_bbox(shuffled_inputs[i], lam) # get random bbox saliencymix_inputs[i, :, bbx1:bbx2, bby1:bby2] = \ shuffled_inputs[i, :, bbx1:bbx2, bby1:bby2] # update the ratio of (area of ori_image on new image) for soft-label lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (ori_inputs.size()[2] * ori_inputs.size()[3])) lam_list.append(lam) labels[i] = labels[i] * lam + shuffled_labels[i] * (1 - lam) long_label = labels.argmax(dim=1) return saliencymix_inputs, labels, long_label class ResizeMix(object): def __init__(self, shuffle_p=1.0, class_num=2, batch_size=4, device='cpu'): """ ResizeMix augmentation arXiv:2012.11101 :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param device: CUDA or CPU """ # ori batch_size=512 self.class_num = class_num self.batch_size = batch_size # calibrate the trigger chance of p, new ratio is the change of operation occur in each batch self.p = shuffle_p self.device = torch.device(device) def __call__(self, inputs, labels, alpha=0.1, beta=0.8, act=True): labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) # duplicate inputs for ori inputs resizemix_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio indices = torch.randperm(self.batch_size, device=self.device) # shuffle indices shuffled_inputs = inputs[indices].to(self.device) shuffled_labels = labels[indices].to(self.device) for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act): # trigger the augmentation operation lam_list.append(-1) continue lam = np.random.uniform(alpha, beta) # lam = 1 - lam bbx1, bby1, bbx2, bby2 = rand_bbox(ori_inputs.size(), lam) # get random bbox # resizer by torchvision torch_resize = Resize([bbx2 - bbx1, bby2 - bby1]) # Tensor -> PIL -> resize -> Tensor re_pil_image = torch_resize(ToPILImage()(shuffled_inputs[i])) resizemix_inputs[i, :, bbx1:bbx2, bby1:bby2] = ToTensor()(re_pil_image) # update the ratio of (area of ori_image on new image) for soft-label lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (ori_inputs.size()[2] * ori_inputs.size()[3])) lam_list.append(lam) labels[i] = labels[i] * lam + shuffled_labels[i] * (1 - lam) long_label = labels.argmax(dim=1) return resizemix_inputs, labels, long_label class FMix(FMixBase): def __init__(self, shuffle_p=1.0, class_num=2, batch_size=4, decay_power=3, alpha=1, size=(32, 32), max_soft=0.0, reformulate=False, device='cpu'): """ FMix augmentation arXiv:2002.12047 :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param batch_size: batch_size of training :param decay_power: decay_power :param alpha: alpha :param size: size of patch :param max_soft: max_soft :param reformulate: reformulate :param device: CUDA or CPU """ # ori batch_size=128 super().__init__(decay_power, alpha, size, max_soft, reformulate) self.class_num = class_num self.batch_size = batch_size self.p = shuffle_p self.device = torch.device(device) def __call__(self, inputs, labels, alpha=1, act=True): # Sample mask and generate random permutation lam, mask = sample_mask(self.alpha, self.decay_power, self.size, self.max_soft, self.reformulate) mask = torch.from_numpy(mask).float().to(self.device) labels = torch.eye(self.class_num).to(self.device)[labels, :] # one-hot hard label ori_inputs = inputs.clone().to(self.device) fmix_inputs = inputs.clone().to(self.device) # duplicate inputs for outputs lam_list = [] # a list to record operating ratio indices = torch.randperm(self.batch_size, device=self.device) # shuffle indices shuffled_inputs = inputs[indices].to(self.device) shuffled_labels = labels[indices].to(self.device) for i in range(self.batch_size): if np.random.randint(0, 101) > 100 * self.p or (not act): # trigger the augmentation operation lam_list.append(-1) continue x1 = mask * ori_inputs[i] x2 = (1 - mask) * shuffled_inputs[i] fmix_inputs[i] = x1 + x2 lam_list.append(lam) labels[i] = labels[i] * lam + shuffled_labels[i] * (1 - lam) long_label = labels.argmax(dim=1) # print('lam:', lam) return fmix_inputs, labels, long_label # CellMix class CellMix(object): def __init__(self, shuffle_p=1.0, class_num=2, strategy='In-place', group_shuffle_size=-1, device='cpu'): """ CellMix augmentation arXiv:2301.11513 :param shuffle_p: chance of trigger augmentation :param class_num: number of classification categories :param strategy: 'In-place' or 'Random' to shuffle the relation patches within the batch :param group_shuffle_size: the size of shuffling group in the batch, -1 to all :param device: CUDA or CPU """ self.p = shuffle_p self.CLS = class_num # classification category number of the task self.device = device self.strategy = strategy # 'In-place' or 'Random' self.group_shuffle_size = group_shuffle_size # -1 for whole batch def __call__(self, inputs, labels, fix_position_ratio=0.5, puzzle_patch_size=32, act=True): """ Fix-position in-place shuffling Perform cross-sample random selection to fix some patches in each image of the batch After selection, the fixed patches are reserved, the rest patches are batch wise in-place shuffled and then regrouped with the fixed patches. cross-sample selection is done by argsort random noise in dim 1 and apply to all image within the batch. in-place batch-wise shuffle operation is done by argsort random noise in dim 0. grouped-in-place batch-wise shuffle operation is done by argsort random noise in the batch dimension :param inputs: input image tensor, size of [B, 3, H, W], :param labels: :param fix_position_ratio: float ratio of the least remaining part of patches :param puzzle_patch_size: int patch size of shuffle :param act: set to be False to force not triggering CellMix in validation, set to True to trigger by chance p output: x, soft_label, long_label x : [B, 3, H, W] re-grouped image after cellmix augmentation soft_label : [B, CLS], soft-label of the class distribution long_label : [B] hard long-label for general discribe """ if np.random.randint(0, 101) > 100 * self.p or (not act): soft_label = torch.eye(self.CLS).to(self.device)[labels, :] # one-hot hard label return inputs, soft_label, labels # Break img into puzzle patches with the size of puzzle_patch_size [B, num_patches, D] inputs = patchify(inputs, puzzle_patch_size) B, num_patches, D = inputs.shape # generate the persudo-mask: in cls dim only the k dim is mask = torch.zeros([B, num_patches, self.CLS], device=inputs.device, requires_grad=False) # no grad # mask of patches: (B, num_patches, cls) (cls)=[0,mask_area,0,....] # transform to persudo-mask B_idx = range(B) mask[B_idx, :, labels] = 1 # num of fix_position puzzle patches len_fix_position = int(num_patches * fix_position_ratio) # create a noise tensor to prepare shuffle idx of puzzle patches noise = torch.rand(1, num_patches, device=self.device) noise = torch.repeat_interleave(noise, repeats=B, dim=0) # based on the batch sequence's shape, the noise tensor get a series idx matrix by sort ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove # sort the idx matrix again, we can obtain the original location idx matrix before assignment ids_restore = torch.argsort(ids_shuffle, dim=1) # keep the first subset ids_fix = ids_shuffle[:, :len_fix_position] # [B,num_patches] -> [B,fix_patches] ids_puzzle = ids_shuffle[:, len_fix_position:] # [B,num_patches] -> [B,puzzle_patches] # set puzzle patch # ids_?.unsqueeze(-1).repeat(1, 1, D) # [B,?_patches] -> [B,?_patches,1] (at each place with the idx of ori patch) -> [B,?_patches,D] # torch.gather to select patche groups x_fixed of [B,fix_patches,D] and x_puzzle of [B,puzzle_patches,D] x_fixed = torch.gather(inputs, dim=1, index=ids_fix.unsqueeze(-1).repeat(1, 1, D)) x_puzzle = torch.gather(inputs, dim=1, index=ids_puzzle.unsqueeze(-1).repeat(1, 1, D)) mask_fixed = torch.gather(mask, dim=1, index=ids_fix.unsqueeze(-1).repeat(1, 1, self.CLS)) mask_puzzle = torch.gather(mask, dim=1, index=ids_puzzle.unsqueeze(-1).repeat(1, 1, self.CLS)) if self.strategy == 'In-place' or self.strategy == 'Random': # the In-place strategy shuffles the relation patches within their location, among the batch index B, num_shuffle_patches, D = x_puzzle.shape # create a noise tensor to prepare shuffle idx of puzzle patches # [B, num_shuffle_patches] noise in [0, 1] noise = torch.rand(B, num_shuffle_patches, device=self.device) if self.group_shuffle_size == -1 or self.group_shuffle_size == B: # CellMix-Split # sort the noise matrix, obtain a index assignment for shuffle, # shuffle dim 0 of entire noise (among all the batch) in_place_shuffle_indices = torch.argsort(noise, dim=0) else: # CellMix-Group assert B > self.group_shuffle_size > 0 and B % self.group_shuffle_size == 0 grouped_indices_list = [] for group_idx in range(B // self.group_shuffle_size): # group the noise by self.group_shuffle_size: [group_shuffle_size,N] grouped_noise = noise[group_idx * self.group_shuffle_size: group_idx * self.group_shuffle_size + self.group_shuffle_size, :] # sort each grouped_noise matrix, obtain a index assignment for shuffle, # now the shuffle dim is 0 (among the batch within the group) grouped_indices = torch.argsort(grouped_noise, dim=0) # put grouped_noise matrix into the list grouped_indices_list.append(grouped_indices + self.group_shuffle_size * group_idx) # stack(cat) the group indices(from list) back to tensor in_place_shuffle_indices = torch.cat(grouped_indices_list, dim=0) # torch.gather to achieve shuffle (taking all the idx base on a shuffled indices) x_puzzle = torch.gather(x_puzzle, dim=0, index=in_place_shuffle_indices.unsqueeze(-1).repeat(1, 1, D)) mask_puzzle = torch.gather(mask_puzzle, dim=0, index=in_place_shuffle_indices.unsqueeze(-1).repeat(1, 1, self.CLS)) else: print('not a valid CellMix strategy') # pack up all puzzle patches inputs = torch.cat([x_fixed, x_puzzle], dim=1) mask = torch.cat([mask_fixed, mask_puzzle], dim=1) # unshuffle to restore the fixed positions inputs = torch.gather(inputs, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, D)) # torch.gather to generate restored binary mask mask = torch.gather(mask, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, self.CLS)) # CellMix random strategy randomly shuffle the image patches (after cellmix in-place shuffle) if self.strategy == 'Random': B, num_patches, D = inputs.shape # create a noise tensor to prepare shuffle idx of puzzle patches noise = torch.rand(B, num_patches, device=self.device) # [num_patches,B] noise in [0, 1] # sort the noise matrix, obtain a index assignment for shuffle, now the shuffle dim is 1 (with the batch) all_shuffle_indices = torch.argsort(noise, dim=1) # ids_shuffle shape of [B,N], in N is idx # torch.gather to shuffle inputs = torch.gather(inputs, dim=1, index=all_shuffle_indices.unsqueeze(-1).repeat(1, 1, D)) # no need to torch the mask, because its patch-wise shuffle within each sample else: # when strategy == 'In-place' pass # unpatchify to obtain puzzle images and their mask inputs = unpatchify(inputs, puzzle_patch_size) # restore to image size：B,3,224,224/ B,3,384,384 # transform soft-mask to soft-label # calaculate a composed label with a conjugate design # [B, num_patches, CLS]->(B, CLS) soft_label = mask.sum(dim=1) # (B, CLS) soft_label = soft_label / num_patches # long_label, as a data-augmentation requirement long_label = soft_label.argmax(dim=1) return inputs, soft_label, long_label # ask func def get_online_augmentation(augmentation_name, p=0.5, class_num=2, batch_size=4, edge_size=224, device='cpu'): """ :param augmentation_name: name of data-augmentation method :param p: chance of triggering :param class_num: classification task num :param batch_size: batch size :param edge_size: edge size of img :param device: cpu or cuda 其中augmentation_name, class_num, batch_size, edge_size必须提供 """ if augmentation_name == 'CellMix-Group': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='In-place', group_shuffle_size=2, device=device) return Augmentation elif augmentation_name == 'CellMix-Group4': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='In-place', group_shuffle_size=4, device=device) return Augmentation elif augmentation_name == 'CellMix-Split': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='In-place', group_shuffle_size=-1, device=device) return Augmentation elif augmentation_name == 'CellMix-Random': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='Random', group_shuffle_size=2, device=device) return Augmentation elif augmentation_name == 'CellMix-Random4': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='Random', group_shuffle_size=4, device=device) return Augmentation elif augmentation_name == 'CellMix-Self': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='Random', group_shuffle_size=1, device=device) return Augmentation elif augmentation_name == 'CellMix-All': Augmentation = CellMix(shuffle_p=p, class_num=class_num, strategy='Random', group_shuffle_size=-1, device=device) return Augmentation elif augmentation_name == 'Cutout': Augmentation = Cutout(alpha=2, shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) return Augmentation elif augmentation_name == 'CutMix': Augmentation = CutMix(alpha=2, shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) return Augmentation elif augmentation_name == 'Mixup': Augmentation = Mixup(alpha=2, shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) return Augmentation elif augmentation_name == 'SaliencyMix': Augmentation = SaliencyMix(alpha=1, shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) # alpha实际为源代码中beta return Augmentation elif augmentation_name == 'ResizeMix': Augmentation = ResizeMix(shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) return Augmentation elif augmentation_name == 'FMix': # FMIX p=1.0 beacuse the chance of trigger is determined inside its own design Augmentation = FMix(shuffle_p=1.0, class_num=class_num, batch_size=batch_size, size=(edge_size, edge_size), device=device) return Augmentation elif augmentation_name == 'PuzzleMix': return None # fixme: all related parts have been taken out seperately # Augmentation = PuzzleMix(alpha=2, shuffle_p=p, class_num=class_num, batch_size=batch_size, device=device) # return Augmentation elif augmentation_name == 'CoMix': # TODO CoMix return None elif augmentation_name == 'RandomMix': # TODO RandomMix return None else: print('no valid counterparts augmentation selected') return None if __name__ == '__main__': ''' Augmentation = get_online_augmentation('CellMix-Split', p=0.5, class_num=2) output, labels, GT_labels = Augmentation(x, label, fix_position_ratio=0.5, puzzle_patch_size=32, act=True) print(labels, GT_labels) ''' x = torch.load("./temp-tensors/warwick.pt") # print(x.shape) label = torch.load("./temp-tensors/warwick_labels.pt") # print(label) # Augmentation = get_online_augmentation('ResizeMix', p=0.5, class_num=2) # output, labels, GT_labels = Augmentation(x, label, act=True) Augmentation = get_online_augmentation('CellMix-Group', p=1, class_num=2) output, labels, GT_labels = Augmentation(x, label, fix_position_ratio=0.5, puzzle_patch_size=32, act=True) print(labels, GT_labels) composed_img = ToPILImage()(output[0]) composed_img.show() composed_img = ToPILImage()(output[1]) composed_img.show() composed_img = ToPILImage()(output[2]) composed_img.show() composed_img = ToPILImage()(output[3]) composed_img.show()