Python torchvision.transforms.transforms.Compose() Examples

def preprocess(image: PIL.Image.Image, image_min_side: float, image_max_side: float) -> Tuple[Tensor, float]:
        # resize according to the rules:
        #   1. scale shorter side to IMAGE_MIN_SIDE
        #   2. after scaling, if longer side > IMAGE_MAX_SIDE, scale longer side to IMAGE_MAX_SIDE
        scale_for_shorter_side = image_min_side / min(image.width, image.height)
        longer_side_after_scaling = max(image.width, image.height) * scale_for_shorter_side
        scale_for_longer_side = (image_max_side / longer_side_after_scaling) if longer_side_after_scaling > image_max_side else 1
        scale = scale_for_shorter_side * scale_for_longer_side

        transform = transforms.Compose([
            transforms.Resize((round(image.height * scale), round(image.width * scale))),  # interpolation `BILINEAR` is applied by default
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = transform(image)

        return image, scale 

def get_transforms(eval=False, aug=None):
    trans = []

    if aug["randcrop"] and not eval:
        trans.append(transforms.RandomCrop(aug["randcrop"]))

    if aug["randcrop"] and eval:
        trans.append(transforms.CenterCrop(aug["randcrop"]))

    if aug["flip"] and not eval:
        trans.append(transforms.RandomHorizontalFlip())

    if aug["grayscale"]:
        trans.append(transforms.Grayscale())
        trans.append(transforms.ToTensor())
        trans.append(transforms.Normalize(mean=aug["bw_mean"], std=aug["bw_std"]))
    elif aug["mean"]:
        trans.append(transforms.ToTensor())
        trans.append(transforms.Normalize(mean=aug["mean"], std=aug["std"]))
    else:
        trans.append(transforms.ToTensor())

    trans = transforms.Compose(trans)
    return trans 

def load_images(img_path):
    # imread from img_path
    img = cv2.imread(img_path)
    img = cv2.resize(img, (224, 224))

    # pytorch must normalize the pic by 
    # mean = [0.485, 0.456, 0.406]
    # std = [0.229, 0.224, 0.225]
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        ])
    
    img = transform(img)
    img.unsqueeze_(0)
    #img_s = img.numpy()
    #img_s = np.transpose(img_s, (1, 2, 0))
    #cv2.imshow("test img", img_s)
    #cv2.waitKey()
    return img 

def preprocess(image: PIL.Image.Image, image_min_side: float, image_max_side: float) -> Tuple[Tensor, float]:
        # resize according to the rules:
        #   1. scale shorter side to IMAGE_MIN_SIDE
        #   2. after scaling, if longer side > IMAGE_MAX_SIDE, scale longer side to IMAGE_MAX_SIDE
        scale_for_shorter_side = image_min_side / min(image.width, image.height)
        longer_side_after_scaling = max(image.width, image.height) * scale_for_shorter_side
        scale_for_longer_side = (image_max_side / longer_side_after_scaling) if longer_side_after_scaling > image_max_side else 1
        scale = scale_for_shorter_side * scale_for_longer_side

        transform = transforms.Compose([
            transforms.Resize((round(image.height * scale), round(image.width * scale))),  # interpolation `BILINEAR` is applied by default
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = transform(image)

        return image, scale 

def prepare_data(images, color_mode='BGR', new_shape=416, color=(127.5, 127.5, 127.5), mode='square'):
    images_ok = np.zeros((images.shape[0], new_shape, new_shape, 3), dtype=images[0].dtype)
    images_tensor = torch.zeros((images.shape[0], 3, new_shape, new_shape), dtype=torch.float32)
    for i in range(len(images)):
        if color_mode == 'BGR':
            images[i] = cv2.cvtColor(images[i], cv2.COLOR_BGR2RGB)
        elif color_mode == 'RGB':
            pass
        else:
            raise NotImplementedError
        images_ok[i], _, _, _ = letterbox(images[i], new_shape, color, mode)

        images_tensor[i] = transforms.Compose([
            transforms.ToPILImage(),
            transforms.ToTensor(),
        ])(images_ok[i])

    return images_tensor 

def get_test_set_lab(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    datasets = [
    EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_TEST_VUE01_30FPS"),
                             transform=transforms.Compose([
                                 RemoveJointsOutsideImgEhpi(image_size),
                                 NormalizeEhpi(image_size)
                             ]), num_joints=num_joints, dataset_part=DatasetPart.TEST),
    EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_TEST_VUE02_30FPS"),
                             transform=transforms.Compose([
                                 RemoveJointsOutsideImgEhpi(image_size),
                                 NormalizeEhpi(image_size)
                             ]), num_joints=num_joints, dataset_part=DatasetPart.TEST),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()
    return ConcatDataset(datasets) 

def get_training_posealgo(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiLSTMDataset] = [
        EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_POSEALGO_30fps"),
                        transform=transforms.Compose([
                            RemoveJointsOutsideImgEhpi(image_size),
                            ScaleEhpi(image_size),
                            TranslateEhpi(image_size),
                            FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                            NormalizeEhpi(image_size)
                        ]), num_joints=num_joints),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_training_set_gt(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiLSTMDataset] = [
        EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_GT_30fps"),
                        transform=transforms.Compose([
                            RemoveJointsOutsideImgEhpi(image_size),
                            ScaleEhpi(image_size),
                            TranslateEhpi(image_size),
                            FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                            NormalizeEhpi(image_size)
                        ]), num_joints=num_joints),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_datasets(initial_pool):
    transform = transforms.Compose(
        [transforms.Resize((224, 224)),
         transforms.RandomHorizontalFlip(),
         transforms.RandomRotation(30),
         transforms.ToTensor(),
         transforms.Normalize(3 * [0.5], 3 * [0.5]), ])
    test_transform = transforms.Compose(
        [
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(3 * [0.5], 3 * [0.5]),
        ]
    )
    # Note: We use the test set here as an example. You should make your own validation set.
    train_ds = datasets.CIFAR10('.', train=True,
                                transform=transform, target_transform=None, download=True)
    test_set = datasets.CIFAR10('.', train=False,
                                transform=test_transform, target_transform=None, download=True)

    active_set = ActiveLearningDataset(train_ds, pool_specifics={'transform': test_transform})

    # We start labeling randomly.
    active_set.label_randomly(initial_pool)
    return active_set, test_set 

def preprocess(image: PIL.Image.Image, image_min_side: float, image_max_side: float) -> Tuple[Tensor, float]:
        # resize according to the rules:
        #   1. scale shorter side to IMAGE_MIN_SIDE
        #   2. after scaling, if longer side > IMAGE_MAX_SIDE, scale longer side to IMAGE_MAX_SIDE
        scale_for_shorter_side = image_min_side / min(image.width, image.height)
        longer_side_after_scaling = max(image.width, image.height) * scale_for_shorter_side
        scale_for_longer_side = (image_max_side / longer_side_after_scaling) if longer_side_after_scaling > image_max_side else 1
        scale = scale_for_shorter_side * scale_for_longer_side

        transform = transforms.Compose([
            transforms.Resize((round(image.height * scale), round(image.width * scale))),  # interpolation `BILINEAR` is applied by default
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = transform(image)

        return image, scale 

def preprocess(self,image: PIL.Image.Image, image_min_side: float, image_max_side: float) -> Tuple[Tensor, float]:
        # resize according to the rules:
        #   1. scale shorter side to IMAGE_MIN_SIDE
        #   2. after scaling, if longer side > IMAGE_MAX_SIDE, scale longer side to IMAGE_MAX_SIDE
        scale_for_shorter_side = image_min_side / min(image.width, image.height)
        longer_side_after_scaling = max(image.width, image.height) * scale_for_shorter_side
        scale_for_longer_side = (image_max_side / longer_side_after_scaling) if longer_side_after_scaling > image_max_side else 1
        scale = scale_for_shorter_side * scale_for_longer_side

        transform = transforms.Compose([
            transforms.Resize((round(image.height * scale), round(image.width * scale))),  # interpolation `BILINEAR` is applied by default
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = transform(image)

        return image, scale 

def pil_to_tensor(img, shape=(64, 64, 3), transform=None):
    """
    Convert PIL image to float tensor

    :param img: PIL image
    :type img: Image.Image
    :param shape: image shape in (H, W, C)
    :type shape: tuple or list
    :param transform: image transform
    :return: tensor
    :rtype: torch.Tensor
    """
    if transform is None:
        transform = transforms.Compose((
            transforms.Resize(shape[0]),
            transforms.ToTensor()
        ))
    return transform(img) 

def get_test_set(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    return EhpiDataset(os.path.join(dataset_path, "2019_03_13_Freilichtmuseum_30FPS"),
                       transform=transforms.Compose([
                           RemoveJointsOutsideImgEhpi(image_size),
                           NormalizeEhpi(image_size)
                       ]), dataset_part=DatasetPart.TEST, num_joints=num_joints) 

def __init__(self, mode, device):

        self.device = device

        if os.path.isdir('/home/scratch/luiraf/work/data/celeba/'):
            data_root = '/home/scratch/luiraf/work/data/celeba/'
        else:
            raise FileNotFoundError('Can\'t find celebrity faces.')

        self.code_root = os.path.dirname(os.path.realpath(__file__))
        self.imgs_root = os.path.join(data_root, 'Img/img_align_celeba/')
        self.imgs_root_preprocessed = os.path.join(data_root, 'Img/img_align_celeba_preprocessed/')
        if not os.path.isdir(self.imgs_root_preprocessed):
            os.mkdir(self.imgs_root_preprocessed)
        self.data_split_file = os.path.join(data_root, 'Eval/list_eval_partition.txt')

        # input: x-y coordinate
        self.num_inputs = 2
        # output: pixel values (RGB)
        self.num_outputs = 3

        # get the labels (train/valid/test)
        train_imgs, valid_imgs, test_imgs = self.get_labels()
        if mode == 'train':
            self.image_files = train_imgs
        elif mode == 'valid':
            self.image_files = valid_imgs
        elif mode == 'test':
            self.image_files = test_imgs
        else:
            raise ValueError

        self.img_size = (32, 32, 3)
        self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
                                             transforms.Resize((self.img_size[0], self.img_size[1]), Image.LANCZOS),
                                             transforms.ToTensor(),
                                             ]) 

def preprocessImage(img, use_color_jitter, image_size_dict, img_norm_info, use_caffe_pretrained_model):
		# calculate target_size and scale_factor, target_size's format is (h, w)
		w_ori, h_ori = img.width, img.height
		if w_ori > h_ori:
			target_size = (image_size_dict.get('SHORT_SIDE'), image_size_dict.get('LONG_SIDE'))
		else:
			target_size = (image_size_dict.get('LONG_SIDE'), image_size_dict.get('SHORT_SIDE'))
		h_t, w_t = target_size
		scale_factor = min(w_t/w_ori, h_t/h_ori)
		target_size = (round(scale_factor*h_ori), round(scale_factor*w_ori))
		# define and do transform
		if use_caffe_pretrained_model:
			means_norm = img_norm_info['caffe'].get('mean_rgb')
			stds_norm = img_norm_info['caffe'].get('std_rgb')
			if use_color_jitter:
				transform = transforms.Compose([transforms.Resize(target_size),
												transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1),
												transforms.ToTensor(),
												transforms.Normalize(mean=means_norm, std=stds_norm)])
			else:
				transform = transforms.Compose([transforms.Resize(target_size),
												transforms.ToTensor(),
												transforms.Normalize(mean=means_norm, std=stds_norm)])
			img = transform(img) * 255
			img = img[(2, 1, 0), :, :]
		else:
			means_norm = img_norm_info['pytorch'].get('mean_rgb')
			stds_norm = img_norm_info['pytorch'].get('std_rgb')
			if use_color_jitter:
				transform = transforms.Compose([transforms.Resize(target_size),
												transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1),
												transforms.ToTensor(),
												transforms.Normalize(mean=means_norm, std=stds_norm)])
			else:
				transform = transforms.Compose([transforms.Resize(target_size),
												transforms.ToTensor(),
												transforms.Normalize(mean=means_norm, std=stds_norm)])
			img = transform(img)
		# return necessary data
		return img, scale_factor, target_size 

def __init__(self, transforms):
        super(Compose, self).__init__(transforms) 

def pathimages_loader(image_paths,size=None,recursive=True,allowed_exts=['jpg', 'jpeg', 'png', 'ppm', 'bmp', 'pgm', 'tif'],shuffle=False,batch_size=32,mean=0.5,std=0.5,transform="default",**loader_args):
    """

    :param image_paths:
    :param size:
    :param recursive:
    :param allowed_exts:
    :param shuffle:
    :param batch_size:
    :param mean:
    :param std:
    :param transform:
    :param loader_args:
    :return:
    """
    if size is not None:
        if not isinstance(size,tuple):
            size = (size,size)

    if transform == "default":
        t = []
        if size is not None:
            t.append(transformations.Resize(size))

        t.append(transformations.ToTensor())

        if mean is not None and std is not None:
            if not isinstance(mean, tuple):
                mean = (mean,)
            if not isinstance(std, tuple):
                std = (std,)
            t.append(transformations.Normalize(mean=mean, std=std))

        trans = transformations.Compose(t)
    else:
        trans = transform

    data = ImagesFromPaths(image_paths,trans,recursive=recursive,allowed_exts=allowed_exts)

    return DataLoader(data,batch_size=batch_size,shuffle=shuffle,**loader_args) 

def fashionmnist_loader(size=None,root="./fashionmnist",train=True,batch_size=32,mean=0.5,std=0.5,transform="default",download=True,target_transform=None,**loader_args):
    """

    :param size:
    :param root:
    :param train:
    :param batch_size:
    :param mean:
    :param std:
    :param transform:
    :param download:
    :param target_transform:
    :param loader_args:
    :return:
    """

    if size is not None:
        if not isinstance(size,tuple):
            size = (size,size)

    if transform == "default":
        t = []
        if size is not None:
            t.append(transformations.Resize(size))

        t.append(transformations.ToTensor())
        if mean is not None and std is not None:
            if not isinstance(mean, tuple):
                mean = (mean,)
            if not isinstance(std, tuple):
                std = (std,)
            t.append(transformations.Normalize(mean=mean, std=std))

        trans = transformations.Compose(t)
    else:
        trans = transform

    data = FashionMNIST(root,train=train,transform=trans,download=download,target_transform=target_transform)

    return DataLoader(data,batch_size=batch_size,shuffle=train,**loader_args) 

def cifar100_loader(size=None,root="./cifar100",train=True,batch_size=32,mean=0.5,std=0.5,transform="default",download=True,target_transform=None,**loader_args):
    """

    :param size:
    :param root:
    :param train:
    :param batch_size:
    :param mean:
    :param std:
    :param transform:
    :param download:
    :param target_transform:
    :param loader_args:
    :return:
    """
    if size is not None:
        if not isinstance(size,tuple):
            size = (size,size)

    if transform == "default":
        t = []
        if size is not None:
            t.append(transformations.Resize(size))

        t.append(transformations.ToTensor())
        if mean is not None and std is not None:
            if not isinstance(mean, tuple):
                mean = (mean,)
            if not isinstance(std, tuple):
                std = (std,)
            t.append(transformations.Normalize(mean=mean, std=std))

        trans = transformations.Compose(t)
    else:
        trans = transform

    data = MNIST(root,train=train,transform=trans,download=download,target_transform=target_transform)

    return DataLoader(data,batch_size=batch_size,shuffle=train,**loader_args) 

def get_train_set(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiDataset] = [
        # Set 1
        EhpiDataset(os.path.join(dataset_path, "ofp_record_2019_03_11_HSRT_30FPS"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints, dataset_part=DatasetPart.TEST),
        # Set 2
        EhpiDataset(os.path.join(dataset_path, "2019_03_13_Freilichtmuseum_30FPS"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints, dataset_part=DatasetPart.TRAIN),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_training_set(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]
    return EhpiDataset(os.path.join(dataset_path, "JHMDB_ITSC-1/"),
                       transform=transforms.Compose([
                           RemoveJointsOutsideImgEhpi(image_size),
                           RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                            probability=0.25),
                           ScaleEhpi(image_size),
                           TranslateEhpi(image_size),
                           FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                           NormalizeEhpi(image_size)
                       ]), num_joints=num_joints) 

def get_sim_pose_algo_only(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiDataset] = [
        EhpiDataset(os.path.join(dataset_path, "ofp_sim_pose_algo_equal_30fps"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
        EhpiDataset(os.path.join(dataset_path, "ofp_from_mocap_pose_algo_30fps"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_sim_gt_only(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiDataset] = [
        EhpiDataset(os.path.join(dataset_path, "ofp_sim_gt_equal_30fps"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
        EhpiDataset(os.path.join(dataset_path, "ofp_from_mocap_gt_30fps"),
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_training_set(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]
    return EhpiDataset(dataset_path,
                       transform=transforms.Compose([
                           RemoveJointsOutsideImgEhpi(image_size),
                           RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                            probability=0.25),
                           ScaleEhpi(image_size),
                           TranslateEhpi(image_size),
                           FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                           NormalizeEhpi(image_size)
                       ]), num_joints=num_joints) 

def get_test_set(image_size: ImageSize):
    num_joints = 15
    return EhpiDataset("/media/disks/beta/datasets/ehpi/2019_03_13_Freilichtmuseum_30FPS",
                       transform=transforms.Compose([
                           RemoveJointsOutsideImgEhpi(image_size),
                           NormalizeEhpi(image_size)
                       ]), dataset_part=DatasetPart.TEST, num_joints=num_joints) 

def get_training_set_both(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiLSTMDataset] = [
        EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_POSEALGO_30fps"),
                        transform=transforms.Compose([
                            RemoveJointsOutsideImgEhpi(image_size),
                            ScaleEhpi(image_size),
                            TranslateEhpi(image_size),
                            FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                            NormalizeEhpi(image_size)
                        ]), num_joints=num_joints),
        EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_03_GT_30fps"),
                        transform=transforms.Compose([
                            RemoveJointsOutsideImgEhpi(image_size),
                            ScaleEhpi(image_size),
                            TranslateEhpi(image_size),
                            FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                            NormalizeEhpi(image_size)
                        ]), num_joints=num_joints),
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def get_test_set_office(dataset_path: str, image_size: ImageSize):
    num_joints = 15
    dataset = EhpiLSTMDataset(os.path.join(dataset_path, "JOURNAL_2019_04_TEST_EVAL2_30FPS"),
                             transform=transforms.Compose([
                                 RemoveJointsOutsideImgEhpi(image_size),
                                 # ScaleEhpi(image_size),
                                 # TranslateEhpi(image_size),
                                 # FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                                 NormalizeEhpi(image_size)
                             ]), num_joints=num_joints, dataset_part=DatasetPart.TEST)
    dataset.print_label_statistics()
    return dataset 

def cmpfacades_loader(size=None,root="./cmpfacades",set="train",batch_size=32,mean=0.5,std=0.5,transform="default",download=True,reverse_mode=False,**loader_args):
    """

    :param size:
    :param root:
    :param set:
    :param batch_size:
    :param mean:
    :param std:
    :param transform:
    :param download:
    :param reverse_mode:
    :param loader_args:
    :return:
    """
    valid_sets = ('train', 'test', 'val')

    if set not in valid_sets: raise ValueError("set {}  is invalid, valid sets include {}".format(set,valid_sets))

    if size is not None:
        if not isinstance(size,tuple):
            size = (size,size)

    if transform == "default":
        t = []
        if size is not None:
            t.append(transformations.Resize(size))

        t.append(transformations.ToTensor())

        if mean is not None and std is not None:
            if not isinstance(mean, tuple):
                mean = (mean,)
            if not isinstance(std, tuple):
                std = (std,)
            t.append(transformations.Normalize(mean=mean, std=std))

        trans = transformations.Compose(t)
    else:
        trans = transform
    data = CMPFacades(root,source_transforms=trans,target_transforms=trans,set=set,download=download,reverse_mode=reverse_mode)
    shuffle_mode = True if set == "train" else False
    return DataLoader(data,batch_size=batch_size,shuffle=shuffle_mode,**loader_args) 

def get_sim(image_size: ImageSize):
    num_joints = 15
    left_indexes: List[int] = [3, 4, 5, 9, 10, 11]
    right_indexes: List[int] = [6, 7, 8, 12, 13, 14]

    datasets: List[EhpiDataset] = [
        EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_sim_pose_algo_equal_30fps",
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
        EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_from_mocap_pose_algo_30fps",
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
        EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_sim_gt_equal_30fps",
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints),
        EhpiDataset("/media/disks/beta/datasets/ehpi/ofp_from_mocap_gt_30fps",
                    transform=transforms.Compose([
                        RemoveJointsOutsideImgEhpi(image_size),
                        RemoveJointsEhpi(indexes_to_remove=foot_indexes, indexes_to_remove_2=knee_indexes,
                                         probability=0.25),
                        ScaleEhpi(image_size),
                        TranslateEhpi(image_size),
                        FlipEhpi(left_indexes=left_indexes, right_indexes=right_indexes),
                        NormalizeEhpi(image_size)
                    ]), num_joints=num_joints)
    ]
    for dataset in datasets:
        dataset.print_label_statistics()

    return ConcatDataset(datasets) 

def handle(self, source, copy_to_local=False, normalize=True,
               split=None, classification_mode=False, **transform_args):
        """

        Args:
            source:
            copy_to_local:
            normalize:
            **transform_args:

        Returns:

        """
        Dataset = self.make_indexing(CelebA)
        data_path = self.get_path(source)

        if copy_to_local:
            data_path = self.copy_to_local_path(data_path)

        if normalize and isinstance(normalize, bool):
            normalize = [(0.5, 0.5, 0.5), (0.5, 0.5, 0.5)]

        if classification_mode:
            train_transform = transforms.Compose([
                transforms.RandomResizedCrop(64),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(*normalize),
            ])
            test_transform = transforms.Compose([
                transforms.Resize(64),
                transforms.CenterCrop(64),
                transforms.ToTensor(),
                transforms.Normalize(*normalize),
            ])
        else:
            train_transform = build_transforms(normalize=normalize,
                                               **transform_args)
            test_transform = train_transform

        if split is None:
            train_set = Dataset(root=data_path, transform=train_transform,
                                download=True)
            test_set = Dataset(root=data_path, transform=test_transform)
        else:
            train_set, test_set = self.make_split(
                data_path, split, Dataset, train_transform, test_transform)
        input_names = ['images', 'labels', 'attributes']

        dim_c, dim_x, dim_y = train_set[0][0].size()
        dim_l = len(train_set.classes)
        dim_a = train_set.attributes[0].shape[0]

        dims = dict(x=dim_x, y=dim_y, c=dim_c, labels=dim_l, attributes=dim_a)
        self.add_dataset('train', train_set)
        self.add_dataset('test', test_set)
        self.set_input_names(input_names)
        self.set_dims(**dims)

        self.set_scale((-1, 1))