class DatasetPerson(Dataset):
    """Dataset containing images from only one person without face detection"""

    def __init__(self, root_dir, transform=None):
        """
        :param root_dir: Directory with the images.
        :param transform: Transformations applied to the images.
        """
        self.transform = transform
        self.root_dir = root_dir
        self.file_names = os.listdir(self.root_dir)
        self.images = []

        # load all images into ram
        for img_name in self.file_names:
            path2img = os.path.join(self.root_dir, img_name)
            img = cv2.imread(path2img)
            if self.transform:
                img = self.transform(img)
            self.images.append(img)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        return self.images[idx]