Python torchvision.transforms.RandomRotation() Examples

def get_dataloaders(data, batch_size=8, study_level=False):
    '''
    Returns dataloader pipeline with data augmentation
    '''
    data_transforms = {
        'train': transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomRotation(10),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 
        ]),
        'valid': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }
    image_datasets = {x: ImageDataset(data[x], transform=data_transforms[x]) for x in data_cat}
    dataloaders = {x: DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in data_cat}
    return dataloaders 

def initialize_dataset(clevr_dir, dictionaries, state_description=True):
    if not state_description:
        train_transforms = transforms.Compose([transforms.Resize((128, 128)),
                                           transforms.Pad(8),
                                           transforms.RandomCrop((128, 128)),
                                           transforms.RandomRotation(2.8),  # .05 rad
                                           transforms.ToTensor()])
        test_transforms = transforms.Compose([transforms.Resize((128, 128)),
                                          transforms.ToTensor()])
                                          
        clevr_dataset_train = ClevrDataset(clevr_dir, True, dictionaries, train_transforms)
        clevr_dataset_test = ClevrDataset(clevr_dir, False, dictionaries, test_transforms)
        
    else:
        clevr_dataset_train = ClevrDatasetStateDescription(clevr_dir, True, dictionaries)
        clevr_dataset_test = ClevrDatasetStateDescription(clevr_dir, False, dictionaries)
    
    return clevr_dataset_train, clevr_dataset_test 

def __call__(self, task_description):
        rotations = {}
        for data_description in task_description:
            c = self.dataset.indices_to_labels[data_description.index]
            if c not in rotations:
                rot = random.choice(self.degrees)
                try:
                    rotations[c] = transforms.Compose([
                        transforms.ToPILImage(),
                        transforms.RandomRotation((rot, rot), fill=(0, )),
                        transforms.ToTensor(),
                    ])
                except Exception:
                    rotations[c] = transforms.Compose([
                        transforms.ToPILImage(),
                        transforms.RandomRotation((rot, rot)),
                        transforms.ToTensor(),
                    ])
            rotation = rotations[c]
            data_description.transforms.append(lambda x: (rotation(x[0]), x[1]))
        return task_description 

def __init__(self, name, size, scale, ratio, colorjitter):
        self.transfs = {
            'val': transforms.Compose([
                transforms.Resize(size),
                transforms.CenterCrop(size=size),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
            'train': transforms.Compose([
                transforms.RandomResizedCrop(size, scale=scale, ratio=ratio),
                transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=colorjitter[0], 
                    contrast=colorjitter[1], 
                    saturation=colorjitter[2]),
                transforms.RandomRotation(degrees=15),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        }[name] 

def __init__(self, transform, mode, select_attrs=[], out_img_size=64, bbox_out_size=32, randomrotate=0, scaleRange=[0.1, 0.9], squareAspectRatio=False, use_celeb=False):
        self.image_path = os.path.join('data','mnist')
        self.mode = mode
        self.iouThresh = 0.5
        self.maxDigits= 1
        self.minDigits = 1
        self.use_celeb = use_celeb
        self.scaleRange = scaleRange
        self.squareAspectRatio = squareAspectRatio
        self.nc = 1 if not self.use_celeb else 3
        transList = [transforms.RandomHorizontalFlip(), transforms.RandomRotation(randomrotate,resample=Image.BICUBIC)]#, transforms.ColorJitter(0.5,0.5,0.5,0.3)
        self.digitTransforms = transforms.Compose(transList)
        self.dataset = MNIST(self.image_path,train=True, transform=self.digitTransforms) if not use_celeb else CelebDataset('./data/celebA/images', './data/celebA/list_attr_celeba.txt', self.digitTransforms, mode)
        self.num_data = len(self.dataset)
        self.metadata = {'images':[]}
        self.catid2attr = {}
        self.out_img_size = out_img_size
        self.bbox_out_size = bbox_out_size
        self.selected_attrs = select_attrs

        print ('Start preprocessing dataset..!')
        self.preprocess()
        print ('Finished preprocessing dataset..!') 

def test_segmentation_pipeline(self):
        class DrawSquare:
            def __init__(self, side):
                self.side = side

            def __call__(self, x, **kwargs):
                x, canvas = x  # x is a [int, ndarray]
                canvas[:self.side, :self.side] = x
                return canvas

        target_trans = BaaLCompose(
            [GetCanvas(), DrawSquare(3), ToPILImage(mode=None), Resize(60, interpolation=0),
             RandomRotation(10, resample=NEAREST, fill=0.0), PILToLongTensor()])
        file_dataset = FileDataset(self.paths, [1] * len(self.paths), self.transform, target_trans)

        x, y = file_dataset[0]
        assert np.allclose(np.unique(y), [0, 1])
        assert y.shape[1:] == x.shape[1:] 

def test_filedataset_segmentation(self):
        target_trans = Compose([default_image_load_fn,
                                Resize(60), RandomRotation(90), ToTensor()])
        file_dataset = FileDataset(self.paths, self.paths, self.transform, target_trans, seed=1337)
        x, y = file_dataset[0]
        assert np.allclose(x.numpy(), y.numpy())
        out1 = list(DataLoader(file_dataset, batch_size=1, num_workers=3, shuffle=False))
        out2 = list(DataLoader(file_dataset, batch_size=1, num_workers=3, shuffle=False))
        assert all([np.allclose(x1.numpy(), x2.numpy())
                    for (x1, _), (x2, _) in zip(out1, out2)])

        file_dataset = FileDataset(self.paths, self.paths, self.transform, target_trans, seed=None)
        x, y = file_dataset[0]
        assert np.allclose(x.numpy(), y.numpy())
        out1 = list(DataLoader(file_dataset, batch_size=1, num_workers=3, shuffle=False))
        out2 = list(DataLoader(file_dataset, batch_size=1, num_workers=3, shuffle=False))
        assert not all([np.allclose(x1.numpy(), x2.numpy())
                        for (x1, _), (x2, _) in zip(out1, out2)]) 

def __init__(self, config=None, batch_size=50, epoch_len = 100, seq_len = 8, transform=None, training_path = '/home/mayank/Desktop/GitRepos/crnn-pytorch_mnist/data/processed/training.pt', ip_channels = 1, size=(20,200), Type='train', profiler = None, target_transform = None):

		self.training_path = training_path
		self.abc = '0123456789'
		self.seq_len = seq_len
		self.epoch_len = epoch_len
		self.transform = transform

		self.train_data, self.train_labels = torch.load(training_path)
		self.num_total = len(self.train_labels)
		self.final_size = size
		self.normal_mean = 7
		self.clip = (1,40)
		self.ip_channels = ip_channels
		self.resized_shape = (*size,ip_channels)

		self.target_aspect_ratio = 10

		self.out_dir = 'out'
		self.rotate = RandomRotation(10)

		self.batch_size = batch_size
		self.encoding_to_char = {1: '0', 2:'1', 3:'2', 4:'3', 5:'4', 6:'5', 7:'6', 8:'7', 9:'8', 10:'9'} 

def __init__(self, patches, use_cache, augment_data):
        super(PatchDataset, self).__init__()
        self.patches = patches
        self.crop = CenterCrop(config.CROP_SIZE)

        if augment_data:
            self.random_transforms = [RandomRotation((90, 90)), RandomVerticalFlip(1.0), RandomHorizontalFlip(1.0),
                                      (lambda x: x)]
            self.get_aug_transform = (lambda: random.sample(self.random_transforms, 1)[0])
        else:
            # Transform does nothing. Not sure if horrible or very elegant...
            self.get_aug_transform = (lambda: (lambda x: x))

        if use_cache:
            self.load_patch = data_manager.load_cached_patch
        else:
            self.load_patch = data_manager.load_patch

        print('Dataset ready with {} tuples.'.format(len(patches))) 

def __init__(self, transform, mode, select_attrs=[], datafile='dataset.json', out_img_size=128, bbox_out_size=64,
                 balance_classes=0, onlyrandBoxes=False, max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1):
        self.image_path = os.path.join('data','belgalogos','images')
        self.transform = transform
        self.mode = mode
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        self.dataset = json.load(open(os.path.join('data','belgalogos',datafile),'r'))
        self.num_data = len(self.dataset['images'])
        self.attr2idx = {}
        self.idx2attr = {}
        self.catid2attr = {}
        self.out_img_size = out_img_size
        self.bbox_out_size = bbox_out_size
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = select_attrs
        self.balance_classes = balance_classes
        self.onlyrandBoxes = onlyrandBoxes
        self.max_object_size = max_object_size
        self.max_with_union= max_with_union
        self.use_gt_mask = use_gt_mask
        self.boxrotate = boxrotate
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])
        if use_gt_mask == 1:
            print ' Not Supported'
            assert(0)

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        self.preprocess()
        print ('Finished preprocessing dataset..!')
        self.imgId2idx = {imid:i for i,imid in enumerate(self.valid_ids)}

        self.num_data = len(self.dataset['images']) 

def __init__(self, transform, mode, select_attrs=[], datafile='dataset.json', out_img_size=128, bbox_out_size=64,
                 balance_classes=0, onlyrandBoxes=False, max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1):
        COCO_classes = ['person' , 'bicycle' , 'car' , 'motorcycle' , 'airplane' , 'bus' , 'train' , 'truck' , 'boat' , 'traffic light' , 'fire hydrant' , 'stop sign' , 'parking meter' , 'bench' , 'bird' , 'cat' , 'dog' , 'horse' , 'sheep' , 'cow' , 'elephant' , 'bear' , 'zebra' , 'giraffe' , 'backpack' , 'umbrella' , 'handbag' , 'tie' , 'suitcase' , 'frisbee' , 'skis' , 'snowboard' , 'sports ball' , 'kite' , 'baseball bat' , 'baseball glove' , 'skateboard' , 'surfboard' , 'tennis racket' , 'bottle' , 'wine glass' , 'cup' , 'fork' , 'knife' , 'spoon' , 'bowl' , 'banana' , 'apple' , 'sandwich' , 'orange' , 'broccoli' , 'carrot' , 'hot dog' , 'pizza' , 'donut' , 'cake' , 'chair' , 'couch' , 'potted plant' , 'bed' , 'dining table' , 'toilet' , 'tv' , 'laptop' , 'mouse' , 'remote' , 'keyboard' , 'cell phone' , 'microwave' , 'oven' , 'toaster' , 'sink' , 'refrigerator' , 'book' , 'clock' , 'vase' , 'scissors' , 'teddy bear' , 'hair drier' , 'toothbrush']
        self.use_cococlass = 1
        self.image_path = os.path.join('data','unrel','images')
        self.transform = transform
        self.mode = mode
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        self.dataset = json.load(open(os.path.join('data','unrel',datafile),'r'))
        self.num_data = len(self.dataset['images'])
        self.attr2idx = {}
        self.idx2attr = {}
        self.catid2attr = {}
        self.out_img_size = out_img_size

        self.bbox_out_size = bbox_out_size
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = COCO_classes if len(select_attrs) == 0 else select_attrs
        self.balance_classes = balance_classes
        self.onlyrandBoxes = onlyrandBoxes
        self.max_object_size = max_object_size
        self.max_with_union= max_with_union
        self.use_gt_mask = 0
        self.boxrotate = boxrotate
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])
        #if use_gt_mask == 1:
        #    print ' Not Supported'
        #    assert(0)

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        self.preprocess()
        print ('Finished preprocessing dataset..!')
        self.imgId2idx = {imid:i for i,imid in enumerate(self.valid_ids)}

        self.num_data = len(self.dataset['images']) 

def __init__(self, transform, mode, select_attrs=[], datafile='dataset.json', out_img_size=128, bbox_out_size=64,
                 balance_classes=0, onlyrandBoxes=False, max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1):
        COCO_classes = ['person' , 'bicycle' , 'car' , 'motorcycle' , 'airplane' , 'bus' , 'train' , 'truck' , 'boat' , 'traffic light' , 'fire hydrant' , 'stop sign' , 'parking meter' , 'bench' , 'bird' , 'cat' , 'dog' , 'horse' , 'sheep' , 'cow' , 'elephant' , 'bear' , 'zebra' , 'giraffe' , 'backpack' , 'umbrella' , 'handbag' , 'tie' , 'suitcase' , 'frisbee' , 'skis' , 'snowboard' , 'sports ball' , 'kite' , 'baseball bat' , 'baseball glove' , 'skateboard' , 'surfboard' , 'tennis racket' , 'bottle' , 'wine glass' , 'cup' , 'fork' , 'knife' , 'spoon' , 'bowl' , 'banana' , 'apple' , 'sandwich' , 'orange' , 'broccoli' , 'carrot' , 'hot dog' , 'pizza' , 'donut' , 'cake' , 'chair' , 'couch' , 'potted plant' , 'bed' , 'dining table' , 'toilet' , 'tv' , 'laptop' , 'mouse' , 'remote' , 'keyboard' , 'cell phone' , 'microwave' , 'oven' , 'toaster' , 'sink' , 'refrigerator' , 'book' , 'clock' , 'vase' , 'scissors' , 'teddy bear' , 'hair drier' , 'toothbrush']
        self.use_cococlass = 1
        self.image_path = os.path.join('data','outofcontext','images')
        self.transform = transform
        self.mode = mode
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        self.dataset = json.load(open(os.path.join('data','outofcontext',datafile),'r'))
        self.num_data = len(self.dataset['images'])
        self.attr2idx = {}
        self.idx2attr = {}
        self.catid2attr = {}
        self.out_img_size = out_img_size

        self.bbox_out_size = bbox_out_size
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = COCO_classes if len(select_attrs) == 0 else select_attrs
        self.balance_classes = balance_classes
        self.onlyrandBoxes = onlyrandBoxes
        self.max_object_size = max_object_size
        self.max_with_union= max_with_union
        self.use_gt_mask = 0
        self.boxrotate = boxrotate
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])
        #if use_gt_mask == 1:
        #    print ' Not Supported'
        #    assert(0)

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        self.preprocess()
        print ('Finished preprocessing dataset..!')
        self.imgId2idx = {imid:i for i,imid in enumerate(self.valid_ids)}

        self.num_data = len(self.dataset['images']) 

def __init__(self, transform, mode, select_attrs=[], datafile='datasetBoxAnn.json', out_img_size=128, bbox_out_size=64,
                 balance_classes=0, onlyrandBoxes=False, max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1):
        self.image_path = os.path.join('data','places2','images')
        self.transform = transform
        self.mode = mode
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        self.filenames = open(os.path.join('data','places2',mode+'_files.txt'),'r').read().splitlines()
        self.num_data = len(self.filenames)
        self.out_img_size = out_img_size
        self.bbox_out_size = bbox_out_size
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = ['background']
        self.onlyrandBoxes = onlyrandBoxes
        self.max_object_size = max_object_size
        self.boxrotate = boxrotate
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        print ('Finished preprocessing dataset..!')

        self.valid_ids  = [int(fname.split('_')[-1].split('.')[0][-8:]) for fname in self.filenames] 

def get_augmentor(is_train, image_size, strong=False):

    augments = []

    if is_train:
        if strong:
            augments.append(transforms.RandomRotation(10))

        augments += [
            transforms.RandomResizedCrop(image_size, interpolation=Image.BILINEAR),
            transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
            transforms.RandomHorizontalFlip()
        ]
    else:
        augments += [
            transforms.Resize(int(image_size / 0.875 + 0.5) if image_size ==
                              224 else image_size, interpolation=Image.BILINEAR),
            transforms.CenterCrop(image_size)
        ]

    augments += [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]

    augmentor = transforms.Compose(augments)
    return augmentor 

def __init__(self, cfg):
        super(SiamFCDataset, self).__init__()
        # pair information
        self.template_size = cfg.SIAMFC.TRAIN.TEMPLATE_SIZE
        self.search_size = cfg.SIAMFC.TRAIN.SEARCH_SIZE
        self.size = (self.search_size - self.template_size) // cfg.SIAMFC.TRAIN.STRIDE + 1   # from cross-correlation

        # aug information
        self.color = cfg.SIAMFC.DATASET.COLOR
        self.flip = cfg.SIAMFC.DATASET.FLIP
        self.rotation = cfg.SIAMFC.DATASET.ROTATION
        self.blur = cfg.SIAMFC.DATASET.BLUR
        self.shift = cfg.SIAMFC.DATASET.SHIFT
        self.scale = cfg.SIAMFC.DATASET.SCALE

        self.transform_extra = transforms.Compose(
            [transforms.ToPILImage(), ] +
            ([transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), ] if self.color > random.random() else [])
            + ([transforms.RandomHorizontalFlip(), ] if self.flip > random.random() else [])
            + ([transforms.RandomRotation(degrees=10), ] if self.rotation > random.random() else [])
        )

        # train data information
        if cfg.SIAMFC.TRAIN.WHICH_USE == 'VID':
            self.anno = cfg.SIAMFC.DATASET.VID.ANNOTATION
            self.num_use = cfg.SIAMFC.TRAIN.PAIRS
            self.root = cfg.SIAMFC.DATASET.VID.PATH
        elif cfg.SIAMFC.TRAIN.WHICH_USE == 'GOT10K':
            self.anno = cfg.SIAMFC.DATASET.GOT10K.ANNOTATION
            self.num_use = cfg.SIAMFC.TRAIN.PAIRS
            self.root = cfg.SIAMFC.DATASET.GOT10K.PATH
        else:
            raise ValueError('not supported training dataset')

        self.labels = json.load(open(self.anno, 'r'))
        self.videos = list(self.labels.keys())
        self.num = len(self.videos)   # video number
        self.frame_range = 100
        self.pick = self._shuffle() 

def setUp(self):
        self.lbls = None
        self.transform = Compose([Resize(60), RandomRotation(90), ToTensor()])
        testtransform = Compose([Resize(32), ToTensor()])
        self.dataset = FileDataset(self.paths, self.lbls, transform=self.transform)
        self.lbls = self.generate_labels(len(self.paths), 10)
        self.dataset = FileDataset(self.paths, self.lbls, transform=self.transform)
        self.active = ActiveLearningDataset(self.dataset,
                                            pool_specifics={'transform': testtransform},
                                            labelled=torch.from_numpy(
                                                (np.array(self.lbls) != -1).astype(np.uint8))) 

def create_loaders(load_random_triplets=False):
    test_dataset_names = copy.copy(dataset_names)
    #test_dataset_names.remove(args.training_set)
    kwargs = {'num_workers': int(args.num_workers), 'pin_memory': args.pin_memory} if args.cuda else {}
    np_reshape64 = lambda x: np.reshape(x, (64, 64, 1))
    transform_test = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform_train = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.RandomRotation(5,PIL.Image.BILINEAR),
            transforms.RandomResizedCrop(32, scale = (0.9,1.0),ratio = (0.9,1.1)),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform = transforms.Compose([
            transforms.Lambda(cv2_scale),
            transforms.Lambda(np_reshape),
            transforms.ToTensor(),
            transforms.Normalize((args.mean_image,), (args.std_image,))])
    if not args.augmentation:
        transform_train = transform
        transform_test = transform


    train_loader = torch.utils.data.DataLoader(
        TotalDatasetsLoader(train=True,
                         load_random_triplets=load_random_triplets,
                         batch_size=args.batch_size,
                         datasets_path=args.dataroot,
                         fliprot=args.fliprot,
                         n_triplets=args.n_triplets,
                         name=args.training_set,
                         download=True,
                         transform=transform_train),
        batch_size=args.batch_size,
        shuffle=False, **kwargs)

    test_loaders = [{'name': name,
                     'dataloader': torch.utils.data.DataLoader(
                         TripletPhotoTour(train=False,
                                          batch_size=args.test_batch_size,
                                          n_triplets = 1000,
                                          root=args.dataroot,
                                          name=name,
                                          download=True,
                                          transform=transform_test),
                         batch_size=args.test_batch_size,
                         shuffle=False, **kwargs)}
                    for name in test_dataset_names]

    return train_loader, test_loaders 

def get_train_loader(data_dir,
                     batch_size,
                     random_seed,
                     shuffle=True,
                     num_workers=4,
                     pin_memory=True):
    """
    Utility function for loading and returning a multi-process
    train iterator over the CIFAR100 dataset.

    If using CUDA, num_workers should be set to 1 and pin_memory to True.

    Args
    ----
    - data_dir: path directory to the dataset.
    - batch_size: how many samples per batch to load.
    - num_workers: number of subprocesses to use when loading the dataset.
    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
      True if using GPU.

    Returns
    -------
    - data_loader: train set iterator.
    """

    # define transforms
    trans = transforms.Compose([
        transforms.RandomCrop(32, padding=4),  # 将图像转化为32 * 32
        transforms.RandomHorizontalFlip(), # 随机水平翻转
        transforms.RandomRotation(degrees=15),  # 随机旋转
        transforms.ToTensor(),  # 将numpy数据类型转化为Tensor
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 归一化
    ])

    # load dataset
    dataset = datasets.CIFAR100(root=data_dir,
                                transform=trans,
                                download=False,
                                train=True)
    if shuffle:
        np.random.seed(random_seed)

    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory,
    )

    return train_loader 

def calc_accuracy(model, input_image_size, use_google_testset=False, testset_path=None, batch_size=32,
                  norm_mean=[0.485, 0.456, 0.406], norm_std=[0.229, 0.224, 0.225]):
    """
    Calculate the mean accuracy of the model on the test test
    :param use_google_testset: If true use the testset derived from google image
    :param testset_path: If None, use a default testset (missing image from the Udacity dataset,
    downloaded from here: http://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz)
    :param batch_size:
    :param model:
    :param input_image_size:
    :param norm_mean:
    :param norm_std:
    :return: the mean accuracy
    """
    if use_google_testset:
        testset_path = "./google_test_data"
        url = 'https://www.dropbox.com/s/3zmf1kq58o909rq/google_test_data.zip?dl=1'
        download_test_set(testset_path, url)
    if testset_path is None:
        testset_path = "./flower_data_orginal_test"
        url = 'https://www.dropbox.com/s/da6ye9genbsdzbq/flower_data_original_test.zip?dl=1'
        download_test_set(testset_path, url)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.eval()
    model.to(device=device)
    with torch.no_grad():
        batch_accuracy = []
        torch.manual_seed(33)
        torch.cuda.manual_seed(33)
        np.random.seed(33)
        random.seed(33)
        torch.backends.cudnn.deterministic = True
        datatransform = transforms.Compose([transforms.RandomRotation(45),
                                            transforms.Resize(input_image_size + 32),
                                            transforms.CenterCrop(input_image_size),
                                            transforms.RandomHorizontalFlip(),
                                            transforms.ToTensor(),
                                            transforms.Normalize(norm_mean, norm_std)])
        image_dataset = datasets.ImageFolder(testset_path, transform=datatransform)
        dataloader = torch.utils.data.DataLoader(image_dataset, batch_size=batch_size, shuffle=True, worker_init_fn=_init_fn)
        for idx, (inputs, labels) in enumerate(dataloader):
            if device == 'cuda':
                inputs, labels = inputs.cuda(), labels.cuda()
            outputs = model.forward(inputs)
            _, predicted = outputs.max(dim=1)
            equals = predicted == labels.data
            print("Batch accuracy (Size {}): {}".format(batch_size, equals.float().mean()))
            batch_accuracy.append(equals.float().mean().cpu().numpy())
        mean_acc = np.mean(batch_accuracy)
        print("Mean accuracy: {}".format(mean_acc))
    return mean_acc 

def loadDataset(dataset, batch_size, train, transform = True):
    oargs = {}
    if dataset in ["MNIST", "CIFAR10", "CIFAR100", "FashionMNIST", "PhotoTour"]:
        oargs['train'] = train
    elif dataset in ["STL10", "SVHN"] :
        oargs['split'] = 'train' if train else 'test'
    elif dataset in ["LSUN"]:
        oargs['classes'] = 'train' if train else 'test'
    elif dataset in ["Imagenet12"]:
        pass
    else:
        raise Exception(dataset + " is not yet supported")

    if dataset in ["MNIST"]:
        transformer = transforms.Compose([ transforms.ToTensor()]
                                         + ([transforms.Normalize((0.1307,), (0.3081,))] if transform else []))
    elif dataset in ["CIFAR10", "CIFAR100"]:
        transformer = transforms.Compose(([ #transforms.RandomCrop(32, padding=4), 
                                            transforms.RandomAffine(0, (0.125, 0.125), resample=PIL.Image.BICUBIC) ,
                                            transforms.RandomHorizontalFlip(), 
                                            #transforms.RandomRotation(15, resample = PIL.Image.BILINEAR) 
                                          ] if train else [])
                                         + [transforms.ToTensor()] 
                                         + ([transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))] if transform else []))
    elif dataset in ["SVHN"]:
        transformer = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (0.2,0.2,0.2))])
    else:
        transformer = transforms.ToTensor()

    if dataset in ["Imagenet12"]:
        # https://github.com/facebook/fb.resnet.torch/blob/master/INSTALL.md#download-the-imagenet-dataset
        train_set = datasets.ImageFolder(
            '../data/Imagenet12/train' if train else '../data/Imagenet12/val',
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                normalize,
            ]))
    else:
        train_set = getattr(datasets, dataset)('../data', download=True, transform=transformer, **oargs)
    return torch.utils.data.DataLoader(
        train_set
        , batch_size=batch_size
        , shuffle=True, 
        **({'num_workers': 1, 'pin_memory': True} if use_cuda else {})) 

def __init__(self, cfg):
        super(SiamRPNDataset, self).__init__()
        # pair information
        self.template_size = cfg.SIAMRPN.TRAIN.TEMPLATE_SIZE
        self.search_size = cfg.SIAMRPN.TRAIN.SEARCH_SIZE
        self.score_size = (self.search_size - self.template_size) // cfg.SIAMRPN.TRAIN.STRIDE + 1  # from cross-correlation

        # anchors information
        self.thr_high = cfg.SIAMRPN.TRAIN.ANCHORS_THR_HIGH
        self.thr_low = cfg.SIAMRPN.TRAIN.ANCHORS_THR_LOW
        self.pos_keep = cfg.SIAMRPN.TRAIN.ANCHORS_POS_KEEP   # kept positive anchors to calc loss
        self.all_keep = cfg.SIAMRPN.TRAIN.ANCHORS_ALL_KEEP   # kept anchors to calc loss
        self.stride = cfg.SIAMRPN.TRAIN.STRIDE
        self.anchor_nums = len(cfg.SIAMRPN.TRAIN.ANCHORS_RATIOS) * len(config.SIAMRPN.TRAIN.ANCHORS_SCALES)
        self._naive_anchors(cfg)   # return self.anchors_naive [anchor_num, 4]
        self._pair_anchors(center=self.search_size//2, score_size=self.score_size)

        # aug information
        self.color = cfg.SIAMRPN.DATASET.COLOR
        self.flip = cfg.SIAMRPN.DATASET.FLIP
        self.rotation = cfg.SIAMRPN.DATASET.ROTATION
        self.blur = cfg.SIAMRPN.DATASET.BLUR
        self.shift = cfg.SIAMRPN.DATASET.SHIFT
        self.scale = cfg.SIAMRPN.DATASET.SCALE
    

        self.transform_extra = transforms.Compose(
            [transforms.ToPILImage(), ] +
            ([transforms.ColorJitter(0.05, 0.05, 0.05, 0.05), ] if self.color > random.random() else [])
            + ([transforms.RandomHorizontalFlip(), ] if self.flip > random.random() else [])
            + ([transforms.RandomRotation(degrees=10), ] if self.rotation > random.random() else [])
        )

        # train data information
        print('train datas: {}'.format(cfg.SIAMRPN.TRAIN.WHICH_USE))
        self.train_datas = []    # all train dataset
        start = 0
        self.num = 0
        for data_name in cfg.SIAMRPN.TRAIN.WHICH_USE:
            dataset = SingleData(cfg, data_name, start)
            self.train_datas.append(dataset)
            start += dataset.num         # real video number
            self.num += dataset.num_use  # the number used for subset shuffle

        # assert abs(self.num - cfg.SIAMRPN.TRAIN.PAIRS) < eps, 'given pairs is not equal to sum of all dataset'

        self._shuffle()
        print(cfg) 

def build_transforms(cfg, is_train=True):
    normalize_transform = T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD)
    if is_train:
        transform_ = T.Compose([
            T.Resize(cfg.INPUT.SIZE_TRAIN),
            T.RandomRotation(cfg.INPUT.RO_DEGREE),
            T.ColorJitter(brightness=cfg.INPUT.BRIGHT_PROB, saturation=cfg.INPUT.SATURA_PROB, contrast=cfg.INPUT.CONTRAST_PROB, hue=cfg.INPUT.HUE_PROB),
            RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN),
            T.Pad(cfg.INPUT.PADDING),
            T.RandomCrop(cfg.INPUT.SIZE_TRAIN),
            T.ToTensor(),
            normalize_transform
        ])
        transform_body = T.Compose([
            T.Resize(cfg.PART.SIZE_BODY),
            T.RandomRotation(cfg.INPUT.RO_DEGREE),
            T.ColorJitter(brightness=cfg.INPUT.BRIGHT_PROB, saturation=cfg.INPUT.SATURA_PROB,
                          contrast=cfg.INPUT.CONTRAST_PROB, hue=cfg.INPUT.HUE_PROB),
            RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN),
            T.Pad(cfg.INPUT.PADDING),
            T.RandomCrop(cfg.PART.SIZE_BODY),
            T.ToTensor(),
            normalize_transform
        ])
        transform_paw = T.Compose([
            T.Resize(cfg.PART.SIZE_PAW),
            T.RandomRotation(cfg.INPUT.RO_DEGREE),
            T.ColorJitter(brightness=cfg.INPUT.BRIGHT_PROB, saturation=cfg.INPUT.SATURA_PROB,
                          contrast=cfg.INPUT.CONTRAST_PROB, hue=cfg.INPUT.HUE_PROB),
            RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN),
            T.Pad(cfg.INPUT.PADDING),
            T.RandomCrop(cfg.PART.SIZE_PAW),
            T.ToTensor(),
            normalize_transform
        ])
        return transform_, transform_body, transform_paw
    else:
        transform = T.Compose([
            T.Resize(cfg.INPUT.SIZE_TEST),
            T.ToTensor(),
            normalize_transform
        ])
        return transform 

def torch_loader(data_path, size):
    if not os.path.exists(data_path+'/train'): download_cifar10(data_path)

    # Data loading code
    traindir = os.path.join(data_path, 'train')
    valdir = os.path.join(data_path, 'test')
    normalize = transforms.Normalize(mean=[0.4914 , 0.48216, 0.44653], std=[0.24703, 0.24349, 0.26159])
    tfms = [transforms.ToTensor(), normalize]

    scale_size = 40
    padding = int((scale_size - size) / 2)
    train_tfms = transforms.Compose([
        pad, # TODO: use `padding` rather than assuming 4
        transforms.RandomCrop(size),
        transforms.ColorJitter(.25,.25,.25),
        transforms.RandomRotation(2),
        transforms.RandomHorizontalFlip(),
    ] + tfms)
    train_dataset = datasets.ImageFolder(traindir, train_tfms)
    train_sampler = (torch.utils.data.distributed.DistributedSampler(train_dataset)
                     if args.distributed else None)
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler)

    val_tfms = transforms.Compose(tfms)
    val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, val_tfms),
        batch_size=args.batch_size*2, shuffle=False,
        num_workers=args.workers, pin_memory=True)


    aug_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, train_tfms),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    train_loader = DataPrefetcher(train_loader)
    val_loader = DataPrefetcher(val_loader)
    aug_loader = DataPrefetcher(aug_loader)
    if args.prof:
        train_loader.stop_after = 200
        val_loader.stop_after = 0

    data = TorchModelData(data_path, args.sz, train_loader, val_loader, aug_loader)
    return data, train_sampler

# Seems to speed up training by ~2% 

def torch_loader(data_path, size):
    if not os.path.exists(data_path+'/train'): download_cifar10(data_path)

    # Data loading code
    traindir = os.path.join(data_path, 'train')
    valdir = os.path.join(data_path, 'test')
    normalize = transforms.Normalize(mean=[0.4914 , 0.48216, 0.44653], std=[0.24703, 0.24349, 0.26159])
    tfms = [transforms.ToTensor(), normalize]

    scale_size = 40
    padding = int((scale_size - size) / 2)
    train_tfms = transforms.Compose([
        pad, # TODO: use `padding` rather than assuming 4
        transforms.RandomCrop(size),
        transforms.ColorJitter(.25,.25,.25),
        transforms.RandomRotation(2),
        transforms.RandomHorizontalFlip(),
    ] + tfms)
    train_dataset = datasets.ImageFolder(traindir, train_tfms)
    train_sampler = (torch.utils.data.distributed.DistributedSampler(train_dataset)
                     if args.distributed else None)
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler)

    val_tfms = transforms.Compose(tfms)
    val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, val_tfms),
        batch_size=args.batch_size*2, shuffle=False,
        num_workers=args.workers, pin_memory=True)


    aug_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, train_tfms),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    train_loader = DataPrefetcher(train_loader)
    val_loader = DataPrefetcher(val_loader)
    aug_loader = DataPrefetcher(aug_loader)
    if args.prof:
        train_loader.stop_after = 200
        val_loader.stop_after = 0

    data = ModelData(data_path, train_loader, val_loader)
    data.sz = args.sz
    data.aug_dl = aug_loader
    if train_sampler is not None: data.trn_sampler = train_sampler
        
    return data

# Seems to speed up training by ~2% 

def create_loaders(load_random_triplets = False):

    test_dataset_names = copy.copy(dataset_names)
    test_dataset_names.remove(args.training_set)

    kwargs = {'num_workers': args.num_workers, 'pin_memory': args.pin_memory} if args.cuda else {}

    np_reshape64 = lambda x: np.reshape(x, (64, 64, 1))
    transform_test = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform_train = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.RandomRotation(5,PIL.Image.BILINEAR),
            transforms.RandomResizedCrop(32, scale = (0.9,1.0),ratio = (0.9,1.1)),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform = transforms.Compose([
            transforms.Lambda(cv2_scale),
            transforms.Lambda(np_reshape),
            transforms.ToTensor(),
            transforms.Normalize((args.mean_image,), (args.std_image,))])
    if not args.augmentation:
        transform_train = transform
        transform_test = transform
    train_loader = torch.utils.data.DataLoader(
            TripletPhotoTour(train=True,
                             load_random_triplets = load_random_triplets,
                             batch_size=args.batch_size,
                             root=args.dataroot,
                             name=args.training_set,
                             download=True,
                             transform=transform_train),
                             batch_size=args.batch_size,
                             shuffle=False, **kwargs)

    test_loaders = [{'name': name,
                     'dataloader': torch.utils.data.DataLoader(
             TripletPhotoTour(train=False,
                     batch_size=args.test_batch_size,
                     root=args.dataroot,
                     name=name,
                     download=True,
                     transform=transform_test),
                        batch_size=args.test_batch_size,
                        shuffle=False, **kwargs)}
                    for name in test_dataset_names]

    return train_loader, test_loaders 

def create_loaders(load_random_triplets=False):
    test_dataset_names = copy.copy(dataset_names)
    #test_dataset_names.remove(args.training_set)
    kwargs = {'num_workers': args.num_workers, 'pin_memory': args.pin_memory} if args.cuda else {}
    np_reshape64 = lambda x: np.reshape(x, (64, 64, 1))
    transform_test = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform_train = transforms.Compose([
            transforms.Lambda(np_reshape64),
            transforms.ToPILImage(),
            transforms.RandomRotation(5,PIL.Image.BILINEAR),
            transforms.RandomResizedCrop(32, scale = (0.9,1.0),ratio = (0.9,1.1)),
            transforms.Resize(32),
            transforms.ToTensor()])
    transform = transforms.Compose([
            transforms.Lambda(cv2_scale),
            transforms.Lambda(np_reshape),
            transforms.ToTensor(),
            transforms.Normalize((args.mean_image,), (args.std_image,))])
    if not args.augmentation:
        transform_train = transform
        transform_test = transform

    train_loader = torch.utils.data.DataLoader(
        TotalDatasetsLoader(train=True,
                         load_random_triplets=load_random_triplets,
                         batch_size=args.batch_size,
                         datasets_path=args.hpatches_split,
                         fliprot=args.fliprot,
                         n_triplets=args.n_triplets,
                         name=args.training_set,
                         download=True,
                         transform=transform_train),
        batch_size=args.batch_size,
        shuffle=False, **kwargs)

    test_loaders = [{'name': name,
                     'dataloader': torch.utils.data.DataLoader(
                         TripletPhotoTour(train=False,
                                          batch_size=args.test_batch_size,
                                          n_triplets = 1000,
                                          root=args.dataroot,
                                          name=name,
                                          download=True,
                                          transform=transform_test),
                         batch_size=args.test_batch_size,
                         shuffle=False, **kwargs)}
                    for name in test_dataset_names]

    return train_loader, test_loaders 

def q_eval(net, dataset, q_idx, flip=False, rotate=False, scale=1):
    # load query image
    q_im = dataset.get_image(q_idx)
    q_size = q_im.size

    # list of transformation lists
    trfs_chains = [[]]
    if rotate:
        eps = 1e-6
        trfs_chains[0] += [RandomRotation((rotate-eps,rotate+eps))]
    if flip:
        trfs_chains[0] += [RandomHorizontalFlip(1)]
    if scale == 0: # AlexNet asks for resized images of 224x224
        edge_list = [224]
        resize_list = [Resize((edge,edge)) for edge in edge_list]
    elif scale == 1:
        edge_list = [800]
        resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
    elif scale == 1.5:
        edge_list = [1200]
        resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
    elif scale == 2: # multiscale
        edge_list = [600,800,1000,1200]
        resize_list = [lambda im: imresize(im, edge) for edge in edge_list]
    else:
        raise ValueError()

    if len(resize_list) == 1:
        trfs_chains[0] += resize_list
    else:
        add_trf(trfs_chains, resize_list )

    # default transformations
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    for chain in trfs_chains:
        chain += [ToTensor(), Normalize(mean, std)]

    net = net.eval()
    q_feat = torch.zeros( (len(trfs_chains), net.out_features) )
    print ('Computing the forward pass and extracting the image representation...')
    for i in range(len(trfs_chains)):
        q_tensor = Compose(trfs_chains[i])(q_im)
        import pdb; pdb.set_trace()  # XXX BREAKPOINT
        q_feat[i] = net.forward(q_tensor.view(1,q_tensor.shape[0],q_tensor.shape[1],q_tensor.shape[2]))
    return F.normalize(q_feat.mean(dim=0), dim=0).detach().numpy() 

def __init__(self, transform, split, select_attrs=[], out_img_size=128, bbox_out_size=64,
                 max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1, square_resize=0) :
        self.image_path = os.path.join('data','ade20k')
        self.transform = transform
        self.split = split
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        datafile = 'train.odgt' if split == 'train' else 'validation.odgt'
        self.datafile = os.path.join('data','ade20k',datafile)
        self.dataset = [json.loads(x.rstrip()) for x in open(self.datafile, 'r')]
        self.num_data = len(self.dataset)
        clsData = open('data/ade20k/object150_info.csv','r').read().splitlines()
        self.clsidx2attr = {i:ln.split(',')[-1] for i, ln in enumerate(clsData[1:])}
        self.clsidx2Stuff = {i:int(ln.split(',')[-2]) for i, ln in enumerate(clsData[1:])}
        self.validCatIds = set([i for i in self.clsidx2Stuff if not self.clsidx2Stuff[i]])
        self.maskSample = 'nonStuff'
        self.out_img_size = out_img_size
        self.square_resize = square_resize
        self.bbox_out_size = bbox_out_size
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = ['background']
        self.max_object_size = max_object_size
        self.max_with_union= max_with_union
        self.use_gt_mask = use_gt_mask
        self.boxrotate = boxrotate
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])
        if use_gt_mask == 1:
            self.mask_transform = transforms.Compose([transforms.Resize(out_img_size if not square_resize else [out_img_size, out_img_size] , interpolation=Image.NEAREST), transforms.CenterCrop(out_img_size)])

        self.valid_ids = []
        for i,img in enumerate(self.dataset):
            imid = int(os.path.basename(img['fpath_img']).split('.')[0].split('_')[-1])
            self.dataset[i]['image_id'] = imid
            self.valid_ids.append(imid)

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        print ('Finished preprocessing dataset..!')
        self.imgId2idx = {imid:i for i,imid in enumerate(self.valid_ids)} 

def __init__(self, transform, mode, select_attrs=[], datafile='datasetBoxAnn.json', out_img_size=128, bbox_out_size=64,
                 balance_classes=0, onlyrandBoxes=False, max_object_size=0., max_with_union=True, use_gt_mask=False,
                 boxrotate=0, n_boxes = 1, square_resize=0, filter_by_mincooccur = -1., only_indiv_occur = 0., augmenter_mode=0):
        self.image_path = os.path.join('data','coco','images')
        self.transform = transform
        self.mode = mode
        self.n_boxes = n_boxes
        self.iouThresh = 0.5
        self.dataset = json.load(open(os.path.join('data','coco',datafile),'r'))
        self.num_data = len(self.dataset['images'])
        self.attr2idx = {}
        self.idx2attr = {}
        self.catid2attr = {}
        self.out_img_size = out_img_size
        self.square_resize = square_resize
        self.bbox_out_size = bbox_out_size
        self.filter_by_mincooccur = filter_by_mincooccur
        self.only_indiv_occur = only_indiv_occur
        #self.selected_attrs = ['person', 'book', 'car', 'bird', 'chair'] if select_attrs== [] else select_attrs
        self.selected_attrs = select_attrs
        if len(select_attrs) == 0:
            self.selected_attrs = [attr['name'] for attr in self.dataset['categories']]
        self.balance_classes = balance_classes
        self.onlyrandBoxes = onlyrandBoxes
        self.max_object_size = max_object_size
        self.max_with_union= max_with_union
        self.use_gt_mask = use_gt_mask
        self.boxrotate = boxrotate
        self.augmenter_mode = augmenter_mode
        if self.boxrotate:
            self.rotateTrans = transforms.Compose([transforms.RandomRotation(boxrotate,resample=Image.NEAREST)])
        if use_gt_mask == 1:
            self.mask_trans =transforms.Compose([transforms.Resize(out_img_size if not square_resize else [out_img_size, out_img_size] , interpolation=Image.NEAREST), transforms.CenterCrop(out_img_size)])
            self.mask_provider = CocoMaskDataset(self.mask_trans, mode, select_attrs=self.selected_attrs, balance_classes=balance_classes)

        self.randHFlip = 'Flip' in transform

        print ('Start preprocessing dataset..!')
        self.preprocess()
        print ('Finished preprocessing dataset..!')
        self.imgId2idx = {imid:i for i,imid in enumerate(self.valid_ids)}

        self.num_data = len(self.dataset['images']) 

def setup_datasets(self):
        """Load the training datasets."""

        train_transform = transforms.Compose(
            [
                transforms.Resize(self.crop_size),
                transforms.RandomRotation(degrees=self.random_angle, resample=Image.BILINEAR),
                transforms.RandomResizedCrop(
                    size=self.crop_size, scale=(1-self.random_scale, 1+self.random_scale), ratio=(1, 1)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]
                )
            ]
        )
        val_transform = transforms.Compose(
            [
                transforms.Resize(self.crop_size),
                transforms.CenterCrop(self.crop_size),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]
                )
            ]
        )

        train_dataset = CocoDatasetPairs(
            root_dir=self.coco_path,
            set_name='train2014',
            transform=train_transform,
            dataset_size_ratio=self.dataset_size_ratio
        )
        train_subset_dataset = Subset(train_dataset, range(0, len(train_dataset), 5*self.dataset_size_ratio))
        val_dataset = CocoDatasetPairs(
            root_dir=self.coco_path,
            set_name='val2014',
            transform=val_transform,
        )

        train_loader = DataLoader(
            train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers
        )
        train_subset_loader = DataLoader(
            train_subset_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers
        )
        return train_loader, train_subset_loader, val_loader