Ejemplo n.º 1
0
Archivo: lfw.py Proyecto: browatbn2/MAD
    def __init__(self, root_dir=cfg.LFW_ROOT, train=True, start=None,
                 max_samples=None, deterministic=True, use_cache=True):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(root_dir, 'crops_tight')
        self.fullsize_img_dir = os.path.join(root_dir, 'images')
        self.feature_dir = os.path.join(root_dir, 'features')

        import glob

        ann = []
        person_dirs = sorted(glob.glob(os.path.join(cfg.LFW_ROOT, 'images', '*')))
        for id, person_dir in enumerate(person_dirs):
            name = os.path.split(person_dir)[1]
            for img_file in sorted(glob.glob(os.path.join(person_dir, '*.jpg'))):
                # create fnames of format 'Aaron_Eckhart/Aaron_Eckhart_0001'
                fname = os.path.join(name, os.path.splitext(os.path.split(img_file)[1])[0])
                ann.append({'fname': fname, 'id': id, 'name': name})

        self.annotations = pd.DataFrame(ann)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        self.transform = ds_utils.build_transform(deterministic=True, color=True)
Ejemplo n.º 2
0
    def __init__(self, root_dir=cfg.AFLW_ROOT, train=True, color=True, start=None,
                 max_samples=None, deterministic=None, use_cache=True,
                 daug=0, return_modified_images=False, test_split='full', align_face_orientation=True,
                 return_landmark_heatmaps=False, landmark_sigma=9, landmark_ids=range(19), **kwargs):

        assert test_split in ['full', 'frontal']
        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.align_face_orientation = align_face_orientation

        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.mode = TRAIN if train else VAL

        self.root_dir = root_dir
        root_dir_local = cfg.AFLW_ROOT_LOCAL
        self.fullsize_img_dir = os.path.join(root_dir, 'data/flickr')
        self.cropped_img_dir = os.path.join(root_dir_local, 'crops')
        self.feature_dir = os.path.join(root_dir_local,  'features')
        self.color = color

        annotation_filename = os.path.join(cfg.AFLW_ROOT_LOCAL, 'alfw.pkl')
        self.annotations_original = pd.read_pickle(annotation_filename)
        print("Number of images: {}".format(len(self.annotations_original)))

        self.frontal_only = test_split == 'frontal'
        self.make_split(train, self.frontal_only)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, True, daug)

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)

        print("Number of images: {}".format(len(self)))
Ejemplo n.º 3
0
Archivo: lfw.py Proyecto: browatbn2/MAD
    def __init__(self, root_dir=cfg.LFW_ROOT, train=True, start=None,
                 max_samples=None, deterministic=True, use_cache=True, view=2):

        assert(view in [1,2])

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else VAL

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(root_dir, 'crops_tight')
        self.fullsize_img_dir = os.path.join(root_dir, 'images')
        self.feature_dir = os.path.join(root_dir, 'features')

        self.default_bbox = [65, 80, 65+100, 80+100]

        pairs_file = 'pairsDevTest.txt' if view == 1 else 'pairs.txt'
        path_annotations = os.path.join(self.root_dir, pairs_file)
        # self.annotations = pd.read_csv(path_annotations)
        self.pairs = []
        with open(path_annotations) as txt_file:
            # num_pairs = int(txt_file.readline()) * 2
            # print(num_pairs)
            for line in txt_file:
                items = line.split()
                if len(items) == 3:
                    pair = (items[0], int(items[1]), items[0], int(items[2]))
                elif len(items) == 4:
                    pair = (items[0], int(items[1]), items[2], int(items[3]))
                else:
                    # print("Invalid line: {}".format(line))
                    continue
                self.pairs.append(pair)

        # assert(num_pairs == len(self.pairs))

        from sklearn.utils import shuffle
        self.pairs = shuffle(self.pairs, random_state=0)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.pairs = self.pairs[st:nd]

        self.transform = ds_utils.build_transform(deterministic, color=True)
Ejemplo n.º 4
0
    def __init__(self, root_dir, fullsize_img_dir, root_dir_local=None, train=True, color=True, start=None,
                 max_samples=None, deterministic=None, use_cache=True, detect_face=False, align_face_orientation=True,
                 return_modified_images=False, return_landmark_heatmaps=True, landmark_sigma=9, landmark_ids=range(68),
                 daug=0, **kwargs):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.train = train
        self.mode = TRAIN if train else VAL
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.start = start
        self.max_samples = max_samples
        self.daug = daug
        self.return_modified_images = return_modified_images

        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.deterministic = deterministic
        if self.deterministic is None:
            self.deterministic = self.mode != TRAIN

        self.fullsize_img_dir = fullsize_img_dir

        self.root_dir = root_dir
        self.root_dir_local = root_dir_local if root_dir_local is not None else self.root_dir

        self.cropped_img_dir = os.path.join(self.root_dir_local, 'crops')
        self.feature_dir = os.path.join(self.root_dir_local,  'features')
        self.color = color

        self.transform = ds_utils.build_transform(self.deterministic, self.color, daug)

        print("Loading annotations... ")
        self.annotations = self.create_annotations()
        print("  Number of images: {}".format(len(self.annotations)))

        self.init()
        self.select_samples()

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)
Ejemplo n.º 5
0
class VggFace2(td.Dataset):

    def __init__(self, root_dir=cfg.VGGFACE2_ROOT, train=True, color=True, start=None,
                 max_samples=None, deterministic=None, min_conf=cfg.MIN_OPENFACE_CONFIDENCE, use_cache=True,
                 crop_source='bb_ground_truth', detect_face=False, align_face_orientation=True,
                 return_landmark_heatmaps=False, return_modified_images=False,
                 daug=0, landmark_sigma=None, landmark_ids=None, **kwargs):

        assert(crop_source in ['bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface'])

        self.mode = TRAIN if train else VAL

        self.face_extractor = FaceExtractor()
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.color = color
        self.crop_source = crop_source
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.root_dir = root_dir
        root_dir_local = cfg.VGGFACE2_ROOT_LOCAL
        split_subfolder = 'train' if train else 'test'
        crop_folder = 'crops'
        if cfg.INPUT_SIZE == 128:
            crop_folder += '_128'
        self.cropped_img_dir = os.path.join(root_dir_local, split_subfolder, crop_folder, crop_source)
        self.fullsize_img_dir = os.path.join(root_dir, split_subfolder, 'imgs')
        self.feature_dir = os.path.join(root_dir_local, split_subfolder, 'features')
        annotation_filename = 'loose_bb_{}.csv'.format(split_subfolder)
        # annotation_filename = 'loose_landmark_{}.csv'.format(split_subfolder)

        # self.path_annotations_mod = os.path.join(root_dir_local, annotation_filename + '.mod_full_of.pkl')
        self.path_annotations_mod = os.path.join(root_dir_local, annotation_filename + '.mod_full.pkl')
        if os.path.isfile(self.path_annotations_mod):
            print('Reading pickle file...')
            self.annotations = pd.read_pickle(self.path_annotations_mod)
            print('done.')
        else:
            print('Reading CSV file...')
            self.annotations = pd.read_csv(os.path.join(self.root_dir, 'bb_landmark', annotation_filename))
            print('done.')

            of_confs, poses, landmarks = [], [], []
            self.annotations = self.annotations[0:4000000]
            self.annotations = self.annotations[self.annotations.H > 80]
            print("Number of images: {}".format(len(self)))

            def get_face_height(lms):
                return lms[8,1] - lms[27,1]

            read_openface_landmarks = True
            if read_openface_landmarks:
                for cnt, filename in enumerate(self.annotations.NAME_ID):
                    filename_noext = os.path.splitext(filename)[0]

                    bb = self.annotations.iloc[cnt][1:5].values
                    expected_face_center = [bb[0] + bb[2] / 2.0, bb[1] + bb[3] / 2.0]

                    conf, lms, pose, num_faces  = ds_utils.read_openface_detection(os.path.join(self.feature_dir, filename_noext),
                                                                       expected_face_center=expected_face_center,
                                                                       use_cache=True, return_num_faces=True)

                    if num_faces > 1:
                        print("Deleting extracted crop for {}...".format(filename))
                        cache_filepath = os.path.join(self.cropped_img_dir, 'tight', filename + '.jpg')
                        if os.path.isfile(cache_filepath):
                            os.remove(cache_filepath)

                    of_confs.append(conf)
                    landmarks.append(lms)
                    poses.append(pose)
                    if (cnt+1) % 10000 == 0:
                        log.info(cnt+1)
                self.annotations['pose'] = poses
                self.annotations['of_conf'] = of_confs
                self.annotations['landmarks_of'] = landmarks

            # assign new continuous ids to persons (0, range(n))
            print("Creating id labels...")
            _ids = self.annotations.NAME_ID
            _ids = _ids.map(lambda x: int(x.split('/')[0][1:]))
            self.annotations['ID'] = _ids

            self.annotations.to_pickle(self.path_annotations_mod)

        min_face_height = 100
        print('Removing faces with height <={:.2f}px...'.format(min_face_height))
        self.annotations = self.annotations[self.annotations.H > min_face_height]
        print("Number of images: {}".format(len(self)))

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color, daug)

        print("Number of images: {}".format(len(self)))
        print("Number of identities: {}".format(self.annotations.ID.nunique()))

    @property
    def labels(self):
        return self.annotations.ID.values

    @property
    def heights(self):
        return self.annotations.H.values

    @property
    def widths(self):
        return self.annotations.W.values

    def get_bounding_box(self, sample):
        bb = sample[1:5].values.copy()

        # convert from x,y,w,h to x1,y1,x2,y2
        bb[2:] += bb[:2]

        # enlarge bounding box
        l,t,r,b = bb
        if t > b:
            t, b = b, t
        h = b-t
        assert(h >= 0)
        t_new, b_new = int(t + 0.05 * h), int(b + 0.1 * h)

        # set width of bbox same as height
        h_new = b_new - t_new
        cx = (r + l) / 2
        l_new, r_new = cx - h_new/2, cx + h_new/2
        # in case right eye is actually left of right eye...
        if l_new > r_new:
            l_new, r_new = r_new, l_new

        bbox = np.array([l_new, t_new, r_new, b_new], dtype=np.float32)
        scalef = cfg.CROP_SIZE / cfg.INPUT_SIZE
        bbox_crop = utils.geometry.scaleBB(bbox, scalef, scalef, typeBB=2)
        return bbox_crop

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

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        filename, id  = sample[0], sample.ID
        bb = None
        landmarks_for_crop = None

        if self.crop_source == 'bb_ground_truth':
            bb = self.get_bounding_box(sample)
            pose = np.zeros(3, dtype=np.float32)
        else:
            of_conf, landmarks_for_crop  = sample.of_conf, sample.landmarks_of
            pose = sample.pose

        try:
            crop, landmarks, pose, cropper = self.face_extractor.get_face(filename+'.jpg', self.fullsize_img_dir,
                                                                          self.cropped_img_dir, landmarks=landmarks_for_crop,
                                                                          bb=bb, pose=pose, use_cache=self.use_cache,
                                                                          detect_face=False, crop_type='tight',
                                                                          aligned=self.align_face_orientation)
        except:
            print(filename)
            raise
            # return self.__getitem__(random.randint(0,len(self)-1))

        try:
            landmarks, _ = cropper.apply_to_landmarks(sample.landmarks)
        except AttributeError:
            landmarks = np.zeros((68,2))

        # vis.show_landmarks(crop, landmarks, title='lms', wait=0, color=(0,0,255))

        cropped_sample = {'image': crop, 'landmarks': landmarks.astype(np.float32), 'pose': pose}

        item = self.transform(cropped_sample)

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        transforms = tf.Compose(transforms)

        result = transforms(item)

        result.update({
            'id': id,
            'fnames': filename,
            'expression': np.array([[0,0,0]], dtype=np.float32),
        })

        if self.return_modified_images:
            mod_transforms = tf.Compose([fp.RandomOcclusion()])
            crop_occ = mod_transforms(item['image'])
            crop_occ = transforms(crop_occ)
            result['image_mod'] = crop_occ

        # add landmark heatmaps if landmarks enabled
        if self.return_landmark_heatmaps:
            result['lm_heatmaps'] = create_landmark_heatmaps(item['landmarks'], self.landmark_sigma, self.landmark_ids)
        return result
Ejemplo n.º 6
0
Archivo: lfw.py Proyecto: browatbn2/MAD
class LFW(td.Dataset):

    def __init__(self, root_dir=cfg.LFW_ROOT, train=True, start=None,
                 max_samples=None, deterministic=True, use_cache=True, view=2):

        assert(view in [1,2])

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else VAL

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(root_dir, 'crops_tight')
        self.fullsize_img_dir = os.path.join(root_dir, 'images')
        self.feature_dir = os.path.join(root_dir, 'features')

        self.default_bbox = [65, 80, 65+100, 80+100]

        pairs_file = 'pairsDevTest.txt' if view == 1 else 'pairs.txt'
        path_annotations = os.path.join(self.root_dir, pairs_file)
        # self.annotations = pd.read_csv(path_annotations)
        self.pairs = []
        with open(path_annotations) as txt_file:
            # num_pairs = int(txt_file.readline()) * 2
            # print(num_pairs)
            for line in txt_file:
                items = line.split()
                if len(items) == 3:
                    pair = (items[0], int(items[1]), items[0], int(items[2]))
                elif len(items) == 4:
                    pair = (items[0], int(items[1]), items[2], int(items[3]))
                else:
                    # print("Invalid line: {}".format(line))
                    continue
                self.pairs.append(pair)

        # assert(num_pairs == len(self.pairs))

        from sklearn.utils import shuffle
        self.pairs = shuffle(self.pairs, random_state=0)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.pairs = self.pairs[st:nd]

        self.transform = ds_utils.build_transform(deterministic, color=True)

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

    def __getitem__(self, idx):
        pair = self.pairs[idx]  # name1, img1, name2, img2

        filepattern = '{}/{}_{:04d}'

        fname1 = filepattern.format(pair[0], pair[0], pair[1])
        fname2 = filepattern.format(pair[2], pair[2], pair[3])

        of_conf1, landmarks, pose = ds_utils.read_openface_detection(os.path.join(self.feature_dir, fname1),
                                                                     expected_face_center=[125,125])

        # if of_conf < 0.025:
        #     landmarks = None
        #     pose = None
        bb = self.default_bbox


        crop, landmarks, pose, cropper = self.face_extractor.get_face(fname1 + '.jpg', self.fullsize_img_dir,
                                                                      self.cropped_img_dir, landmarks=landmarks,
                                                                      pose=pose, use_cache=self.use_cache,
                                                                      bb=bb,
                                                                      detect_face=False, crop_type='tight',
                                                                      aligned=True)

        transformed_crop1 = self.transform(crop)

        of_conf2, landmarks, pose = ds_utils.read_openface_detection(os.path.join(self.feature_dir, fname2),
                                                                     expected_face_center=[125,125])
        # if of_conf < 0.025:
        #     landmarks = None
        #     pose = None
        crop, landmarks, pose, cropper = self.face_extractor.get_face(fname2 + '.jpg', self.fullsize_img_dir,
                                                                      self.cropped_img_dir, landmarks=landmarks,
                                                                      pose=pose, use_cache=self.use_cache,
                                                                      bb=bb,
                                                                      detect_face=False, crop_type='tight',
                                                                      aligned=True)
            # import matplotlib.pyplot as plt
            # plt.imshow(crop)
            # plt.show()
        transformed_crop2 = self.transform(crop)

        return transformed_crop1, transformed_crop2, pair[0], pair[2], pair[0]==pair[2], float(of_conf1), float(of_conf2)


    def get_face(self, filename, bb, landmarks=None, size=(cfg.CROP_SIZE, cfg.CROP_SIZE)):
        # Load image from dataset
        img_path = os.path.join(self.fullsize_img_dir, filename + '.jpg')
        img = io.imread(img_path)
        if img is None:
            raise IOError("\tError: Could not load image {}!".format(img_path))

        #
        # Crop face using landmarks or bounding box
        #

        def crop_by_bb(img, bb):
            x, y, w, h = bb
            x, y, x2, y2 = max(0, x), max(0, y), min(img.shape[1], x + w), min(img.shape[0], y + h)
            return img[y:y2, x:x2]

        def crop_by_lm(img, landmarks):
            return face_processing.crop_bump(img, landmarks, output_size=size)

        # print(filename)

        # load landmarks extracted with OpenFace2
        pose = np.zeros(3, dtype=np.float32)
        lmFilepath = os.path.join(self.feature_dir, filename + '.csv')
        try:
            features = pd.read_csv(lmFilepath, skipinitialspace=True)
            features.sort_values('confidence', ascending=False)
            if features.confidence[0] > 0.0:
                landmarks_x = features.as_matrix(columns=['x_{}'.format(i) for i in range(68)])[0]
                landmarks_y = features.as_matrix(columns=['y_{}'.format(i) for i in range(68)])[0]
                landmarks = np.vstack((landmarks_x, landmarks_y)).T
                pitch = features.pose_Rx.values[0]
                yaw = features.pose_Ry.values[0]
                roll = features.pose_Rz.values[0]
                pose = np.array((pitch, yaw, roll), dtype=np.float32)
        except IOError:
            # raise IOError("\tError: Could not load landmarks from file {}!".format(lmFilepath))
            print("\tError: Could not load landmarks from file {}!".format(lmFilepath))

        if landmarks is not None:
            crop, landmarks = crop_by_lm(img, landmarks)
        else:
            crop, landmarks = crop_by_bb(img, bb), np.zeros((68, 2), dtype=np.float32)

        return cv2.resize(crop, size, interpolation=cv2.INTER_CUBIC), pose, landmarks
Ejemplo n.º 7
0
Archivo: lfw.py Proyecto: browatbn2/MAD
class LFWImages(td.Dataset):

    def __init__(self, root_dir=cfg.LFW_ROOT, train=True, start=None,
                 max_samples=None, deterministic=True, use_cache=True):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(root_dir, 'crops_tight')
        self.fullsize_img_dir = os.path.join(root_dir, 'images')
        self.feature_dir = os.path.join(root_dir, 'features')

        import glob

        ann = []
        person_dirs = sorted(glob.glob(os.path.join(cfg.LFW_ROOT, 'images', '*')))
        for id, person_dir in enumerate(person_dirs):
            name = os.path.split(person_dir)[1]
            for img_file in sorted(glob.glob(os.path.join(person_dir, '*.jpg'))):
                # create fnames of format 'Aaron_Eckhart/Aaron_Eckhart_0001'
                fname = os.path.join(name, os.path.splitext(os.path.split(img_file)[1])[0])
                ann.append({'fname': fname, 'id': id, 'name': name})

        self.annotations = pd.DataFrame(ann)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        self.transform = ds_utils.build_transform(deterministic=True, color=True)

    @property
    def labels(self):
        return self.annotations.id.values

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

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        filename = sample.fname
        id = sample.id
        of_conf, landmarks, pose = ds_utils.read_openface_detection(os.path.join(self.feature_dir, filename))
        # if of_conf < 0.8:
        #     return self.__getitem__((idx+1) % len(self))

        try:
            # crop, landmarks, pose = ds_utils.get_face(filename+'.jpg', self.fullsize_img_dir, self.cropped_img_dir,
            #                                           landmarks, pose, use_cache=False)
            crop, landmarks, pose, cropper = self.face_extractor.get_face(filename + '.jpg', self.fullsize_img_dir,
                                                                          self.cropped_img_dir, landmarks=landmarks,
                                                                          pose=pose, use_cache=self.use_cache,
                                                                          detect_face=False, crop_type='tight',
                                                                          aligned=True)
        except:
            print(filename)
            return self.__getitem__((idx+1) % len(self))

        # vis.show_landmarks(crop, landmarks, pose=pose, title='lms', wait=10, color=(0,0,255))

        transformed_crop = self.transform(crop)

        landmarks[..., 0] -= int((crop.shape[0] - transformed_crop.shape[1]) / 2)
        landmarks[..., 1] -= int((crop.shape[1] - transformed_crop.shape[2]) / 2)

        item = {
            'image': transformed_crop,
            'id': id,
            'fnames': filename,
            'pose': pose,
            'landmarks': landmarks,
        }
        return item
Ejemplo n.º 8
0
    def __init__(self,
                 root_dir=cfg.AFFECTNET_ROOT,
                 train=True,
                 transform=None,
                 crop_type='tight',
                 color=True,
                 start=None,
                 max_samples=None,
                 outlier_threshold=None,
                 deterministic=None,
                 use_cache=True,
                 detect_face=False,
                 align_face_orientation=False,
                 min_conf=cfg.MIN_OPENFACE_CONFIDENCE,
                 daug=0,
                 return_landmark_heatmaps=False,
                 landmark_sigma=9,
                 landmark_ids=range(68),
                 return_modified_images=False,
                 crop_source='lm_openface',
                 **kwargs):
        assert (crop_type in ['fullsize', 'tight', 'loose'])
        assert (crop_source in [
            'bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface'
        ])

        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else VAL

        self.crop_source = crop_source
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.start = start
        self.max_samples = max_samples

        self.root_dir = root_dir
        self.crop_type = crop_type
        self.color = color
        self.outlier_threshold = outlier_threshold
        self.transform = transform
        self.fullsize_img_dir = os.path.join(self.root_dir,
                                             'cropped_Annotated')
        self.cropped_img_dir = os.path.join(self.root_dir, 'crops',
                                            crop_source)
        self.feature_dir = os.path.join(self.root_dir, 'features')

        annotation_filename = 'training' if train else 'validation'
        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod):
            print('Reading pickle file...')
            self._annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading CSV file...')
            self._annotations = pd.read_csv(
                os.path.join(root_dir, annotation_filename + '.csv'))
            print('done.')

            # drop non-faces
            self._annotations = self._annotations[
                self._annotations.expression < 8]

            # Samples in annotation file are somewhat clustered by expression.
            # Shuffle to create a more even distribution.
            # NOTE: deterministic, always creates the same order
            if train:
                from sklearn.utils import shuffle
                self._annotations = shuffle(self._annotations, random_state=2)

                # remove samples with inconsistent expression<->valence/arousal values
                self._remove_outliers()

            poses = []
            confs = []
            landmarks = []
            for cnt, filename in enumerate(
                    self._annotations.subDirectory_filePath):
                if cnt % 1000 == 0:
                    print(cnt)
                filename_noext = os.path.splitext(filename)[0]
                conf, lms, pose = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename_noext))
                poses.append(pose)
                confs.append(conf)
                landmarks.append(lms)
            self._annotations['pose'] = poses
            self._annotations['conf'] = confs
            self._annotations['landmarks_of'] = landmarks
            # self.annotations.to_csv(path_annotations_mod, index=False)
            self._annotations.to_pickle(path_annotations_mod)

        poses = np.abs(np.stack(self._annotations.pose.values))

        only_good_image_for_training = True
        if train and only_good_image_for_training:
            print(len(self._annotations))

            min_rot_deg = 30
            max_rot_deg = 90
            # print('Limiting rotation to +-[{}-{}] degrees...'.format(min_rot_deg, max_rot_deg))
            # self._annotations = self._annotations[(poses[:, 0] < np.deg2rad(max_rot_deg)) &
            #                                       (poses[:, 1] < np.deg2rad(max_rot_deg)) &
            #                                       (poses[:, 2] < np.deg2rad(max_rot_deg))]
            # self._annotations = self._annotations[(np.deg2rad(min_rot_deg) < poses[:, 0]) |
            #                                       (np.deg2rad(min_rot_deg) < poses[:, 1])]
            # self._annotations = self._annotations[np.deg2rad(min_rot_deg) < poses[:, 1] ]

            print(len(self._annotations))

            # print('Removing OpenFace confs <={:.2f}...'.format(min_conf))
            # self._annotations = self._annotations[self._annotations.conf > cfg.MIN_OPENFACE_CONFIDENCE]
            # print(len(self._annotations))

            # select by Valence/Arousal
            # min_arousal = 0.0
            # print('Removing arousal <={:.2f}...'.format(min_arousal))
            # self._annotations = self._annotations[self._annotations.arousal > min_arousal]
            # print(len(self._annotations))

        # There is (at least) one missing image in the dataset. Remove by checking face width:
        self._annotations = self._annotations[self._annotations.face_width > 0]

        # self._annotations_balanced = self._annotations
        # self.filter_labels(label_dict_exclude={'expression': 0})
        # self.filter_labels(label_dict_exclude={'expression': 1})
        # self._annotations = self._annotations[self._annotations.arousal > 0.2]

        self.rebalance_classes()

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color,
                                                  daug)

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor()]
        transforms += [fp.Normalize([0.518, 0.418, 0.361],
                                    [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)
Ejemplo n.º 9
0
    def __init__(self,
                 root_dir=cfg.VOXCELEB_ROOT,
                 train=True,
                 start=None,
                 max_samples=None,
                 deterministic=True,
                 with_bumps=False,
                 min_of_conf=0.3,
                 min_face_height=100,
                 use_cache=True,
                 **kwargs):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(cfg.VOXCELEB_ROOT_LOCAL, 'crops')
        self.fullsize_img_dir = os.path.join(
            root_dir, 'frames/unzippedIntervalFaces/data')
        self.feature_dir = os.path.join(root_dir,
                                        'features/unzippedIntervalFaces/data')
        self.npfeature_dir = os.path.join(
            cfg.VOXCELEB_ROOT_LOCAL, 'features/unzippedIntervalFaces/data')
        self.train = train
        self.with_bumps = with_bumps

        annotation_filename = 'dev' if train else 'test'
        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod) and False:
            self.annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading CSV file...')
            self.annotations = pd.read_csv(
                os.path.join(root_dir, annotation_filename + '.csv'))
            print('done.')

            # self.annotations['of_conf'] = -1
            # self.annotations['landmarks'] = ''
            # self.annotations['pose'] = ''
            # of_confs, poses, landmarks = [], [], []
            #
            #
            #
            # # for cnt, filename in enumerate(self.annotations.fname):
            # for cnt, idx in enumerate(self.annotations.index):
            #     filename = self.annotations.iloc[idx].fname
            #     filename_noext = os.path.splitext(filename)[0]
            #     of_conf, lms, pose = ds_utils.read_openface_detection(os.path.join(self.feature_dir, filename_noext))
            #     str_landmarks = encode_landmarks(lms)
            #     of_confs.append(of_conf)
            #     # poses.append(pose)
            #     landmarks.append(lms)
            #     self.annotations.loc[idx, 'of_conf'] = of_conf
            #     self.annotations.loc[idx, 'landmarks'] = str_landmarks
            #     self.annotations.loc[idx, 'pose'] = encode_landmarks(pose)
            #     if (cnt+1) % 100 == 0:
            #         print(cnt+1)
            #     if (cnt+1) % 1000 == 0:
            #         print('saving annotations...')
            #         self.annotations.to_pickle(path_annotations_mod)
            # # self.annotations.to_csv(path_annotations_mod, index=False)
            # self.annotations.to_pickle(path_annotations_mod)

        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.lms.pkl')
        lm_annots = pd.read_pickle(os.path.join(root_dir,
                                                path_annotations_mod))

        t = time.time()
        self.annotations = pd.merge(self.annotations,
                                    lm_annots,
                                    on='fname',
                                    how='inner')
        print("Time merge: {:.2f}".format(time.time() - t))

        t = time.time()
        self.annotations['vid'] = self.annotations.fname.map(
            lambda x: x.split('/')[2])
        self.annotations['id'] = self.annotations.uid.map(lambda x: int(x[2:]))
        print("Time vid/id labels: {:.2f}".format(time.time() - t))

        print("Num. faces: {}".format(len(self.annotations)))
        print("Num. ids  : {}".format(self.annotations.id.nunique()))

        # drop bad face detections
        print("Removing faces with conf < {}".format(min_of_conf))
        self.annotations = self.annotations[
            self.annotations.of_conf >= min_of_conf]
        print("Num. faces: {}".format(len(self.annotations)))

        # drop small faces
        print("Removing faces with height < {}px".format(min_face_height))
        self.annotations = self.annotations[
            self.annotations.face_height >= min_face_height]
        print("Num. faces: {}".format(len(self.annotations)))

        fr = 0
        prev_vid = -1
        frame_nums = []
        for n, id in enumerate(self.annotations.vid.values):
            fr += 1
            if id != prev_vid:
                prev_vid = id
                fr = 0
            frame_nums.append(fr)
        self.annotations['FRAME'] = frame_nums

        self.max_frames_per_video = 200
        self.frame_interval = 3
        print('Limiting videos in VoxCeleb to {} frames...'.format(
            self.max_frames_per_video))
        self.annotations = self.annotations[self.annotations.FRAME %
                                            self.frame_interval == 0]
        self.annotations = self.annotations[
            self.annotations.FRAME < self.max_frames_per_video *
            self.frame_interval]
        print("Num. faces: {}".format(len(self.annotations)))

        # limit number of samples
        st, nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st + max_samples
        self.annotations = self.annotations[st:nd]

        self.transform = ds_utils.build_transform(deterministic=True,
                                                  color=True)
Ejemplo n.º 10
0
class VoxCeleb(td.Dataset):
    def __init__(self,
                 root_dir=cfg.VOXCELEB_ROOT,
                 train=True,
                 start=None,
                 max_samples=None,
                 deterministic=True,
                 with_bumps=False,
                 min_of_conf=0.3,
                 min_face_height=100,
                 use_cache=True,
                 **kwargs):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.root_dir = root_dir
        self.cropped_img_dir = os.path.join(cfg.VOXCELEB_ROOT_LOCAL, 'crops')
        self.fullsize_img_dir = os.path.join(
            root_dir, 'frames/unzippedIntervalFaces/data')
        self.feature_dir = os.path.join(root_dir,
                                        'features/unzippedIntervalFaces/data')
        self.npfeature_dir = os.path.join(
            cfg.VOXCELEB_ROOT_LOCAL, 'features/unzippedIntervalFaces/data')
        self.train = train
        self.with_bumps = with_bumps

        annotation_filename = 'dev' if train else 'test'
        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod) and False:
            self.annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading CSV file...')
            self.annotations = pd.read_csv(
                os.path.join(root_dir, annotation_filename + '.csv'))
            print('done.')

            # self.annotations['of_conf'] = -1
            # self.annotations['landmarks'] = ''
            # self.annotations['pose'] = ''
            # of_confs, poses, landmarks = [], [], []
            #
            #
            #
            # # for cnt, filename in enumerate(self.annotations.fname):
            # for cnt, idx in enumerate(self.annotations.index):
            #     filename = self.annotations.iloc[idx].fname
            #     filename_noext = os.path.splitext(filename)[0]
            #     of_conf, lms, pose = ds_utils.read_openface_detection(os.path.join(self.feature_dir, filename_noext))
            #     str_landmarks = encode_landmarks(lms)
            #     of_confs.append(of_conf)
            #     # poses.append(pose)
            #     landmarks.append(lms)
            #     self.annotations.loc[idx, 'of_conf'] = of_conf
            #     self.annotations.loc[idx, 'landmarks'] = str_landmarks
            #     self.annotations.loc[idx, 'pose'] = encode_landmarks(pose)
            #     if (cnt+1) % 100 == 0:
            #         print(cnt+1)
            #     if (cnt+1) % 1000 == 0:
            #         print('saving annotations...')
            #         self.annotations.to_pickle(path_annotations_mod)
            # # self.annotations.to_csv(path_annotations_mod, index=False)
            # self.annotations.to_pickle(path_annotations_mod)

        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.lms.pkl')
        lm_annots = pd.read_pickle(os.path.join(root_dir,
                                                path_annotations_mod))

        t = time.time()
        self.annotations = pd.merge(self.annotations,
                                    lm_annots,
                                    on='fname',
                                    how='inner')
        print("Time merge: {:.2f}".format(time.time() - t))

        t = time.time()
        self.annotations['vid'] = self.annotations.fname.map(
            lambda x: x.split('/')[2])
        self.annotations['id'] = self.annotations.uid.map(lambda x: int(x[2:]))
        print("Time vid/id labels: {:.2f}".format(time.time() - t))

        print("Num. faces: {}".format(len(self.annotations)))
        print("Num. ids  : {}".format(self.annotations.id.nunique()))

        # drop bad face detections
        print("Removing faces with conf < {}".format(min_of_conf))
        self.annotations = self.annotations[
            self.annotations.of_conf >= min_of_conf]
        print("Num. faces: {}".format(len(self.annotations)))

        # drop small faces
        print("Removing faces with height < {}px".format(min_face_height))
        self.annotations = self.annotations[
            self.annotations.face_height >= min_face_height]
        print("Num. faces: {}".format(len(self.annotations)))

        fr = 0
        prev_vid = -1
        frame_nums = []
        for n, id in enumerate(self.annotations.vid.values):
            fr += 1
            if id != prev_vid:
                prev_vid = id
                fr = 0
            frame_nums.append(fr)
        self.annotations['FRAME'] = frame_nums

        self.max_frames_per_video = 200
        self.frame_interval = 3
        print('Limiting videos in VoxCeleb to {} frames...'.format(
            self.max_frames_per_video))
        self.annotations = self.annotations[self.annotations.FRAME %
                                            self.frame_interval == 0]
        self.annotations = self.annotations[
            self.annotations.FRAME < self.max_frames_per_video *
            self.frame_interval]
        print("Num. faces: {}".format(len(self.annotations)))

        # limit number of samples
        st, nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st + max_samples
        self.annotations = self.annotations[st:nd]

        self.transform = ds_utils.build_transform(deterministic=True,
                                                  color=True)

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Train: {}\n'.format(self.train)
        fmt_str += '    Root Location: {}\n'.format(self.root_dir)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(
            tmp,
            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        fmt_str += self._stats_repr()
        return fmt_str

    def _stats_repr(self):
        fmt_str = "     Number of identities: {}\n".format(
            self.annotations.id.nunique())
        fmt_str += "     Number of videos:     {}\n".format(
            self.annotations.vid.nunique())
        fmt_str += "     Frame inverval:       {}\n".format(
            self.frame_interval)
        fmt_str += "     Max frames per vid:   {}\n".format(
            self.max_frames_per_video)
        return fmt_str

    @property
    def labels(self):
        return self.annotations.id.values

    def get_landmarks(self, sample):
        landmarks = np.array([sample.landmarks_x, sample.landmarks_y],
                             dtype=np.float32).T
        # return face_processing.scale_landmarks_to_crop(landmarks, output_size=(cfg.CROP_SIZE, cfg.CROP_SIZE))
        return landmarks

    @property
    def vids(self):
        return self.annotations.vid.values

    def show_landmarks(self, img, landmarks, title='landmarks'):
        for lm in landmarks:
            lm_x, lm_y = lm[0], lm[1]
            cv2.circle(img, (int(lm_x), int(lm_y)), 2, (0, 0, 255), -1)
        cv2.imshow(title, cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        cv2.waitKey(10)

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

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        filename, id, vid = sample.fname, sample.id, sample.vid

        pose = np.array((sample.pose_pitch, sample.pose_yaw, sample.pose_roll),
                        dtype=np.float32)
        landmarks = self.get_landmarks(sample)

        t = time.time()
        crop, landmarks, pose, cropper = self.face_extractor.get_face(
            filename + '.jpg',
            self.fullsize_img_dir,
            self.cropped_img_dir,
            landmarks=landmarks,
            pose=pose,
            use_cache=self.use_cache,
            detect_face=False,
            crop_type='tight',
            aligned=True)
        # self.show_landmarks(crop, landmarks, 'imgs')
        # crop, pose, landmarks,of_conf_seq = self.get_face(fname, landmarks=landmarks, use_cache=True, from_sequence=True)
        # landmarks = face_processing.scale_landmarks_to_crop(landmarks, output_size=(cfg.CROP_SIZE, cfg.CROP_SIZE))
        # self.show_landmarks(crop, landmarks, 'sequence')
        # print(of_conf, of_conf_seq)
        # cv2.waitKey()

        cropped_sample = {'image': crop, 'landmarks': landmarks, 'pose': pose}
        item = self.transform(cropped_sample)

        # face_mask = face_processing.get_face_mask(landmarks, crop.shape)
        # transformed_face_mask = face_processing.CenterCrop(cfg.INPUT_SIZE)(face_mask)

        item.update({
            'id': id,
            'fnames': filename,
            # 'face_mask': transformed_face_mask,
            'expression': np.array([[-1, 0, 0]], np.float32),
            'vid': vid
        })

        if self.with_bumps:
            H = np.eye(3, 3)
            step = 1
            next_id = max(0, min(len(self) - 1, idx + step))
            if self.annotations.iloc[next_id].vid != vid:
                next_id = max(0, min(len(self) - 1, idx - step))
            if self.annotations.iloc[max(0, idx - step)].vid != vid:
                # fallback to single image
                crop_next = crop
                landmarks_next = landmarks
            else:
                sample_next = self.annotations.iloc[next_id]
                pose_next = np.array(
                    (sample_next.pose_pitch, sample_next.pose_yaw,
                     sample_next.pose_roll),
                    dtype=np.float32)
                of_conf = sample_next.of_conf
                crop_next, landmarks_next = self.get_face(
                    sample_next.fname,
                    landmarks=self.get_landmarks(sample_next),
                    use_cache=True)

                pose_diff = np.abs(pose - pose_next)
                if np.any(pose_diff > np.deg2rad(7)) or of_conf < 0.9:
                    # print(np.rad2deg(pose_diff))
                    crop_next = crop
                    landmarks_next = landmarks
                else:
                    # calculate homograpy to project next images onto current image
                    H = cv2.findHomography(landmarks_next, landmarks)[0]

            bumps = face_processing.calc_face_bumps(crop, crop_next, landmarks,
                                                    landmarks_next, H)
            transformed_bumps = face_processing.CenterCrop(
                cfg.INPUT_SIZE)(bumps)
            # transformed_bumps = self.transform(bumps)

            # bumps = face_processing.calc_face_bumps(crop, crop_next, landmarks, landmarks_next)
            # transformed_crop_next = self.transform(crop_next)
            # return transformed_crop, id, pose, landmarks, emotion, vid, fname, transformed_crop_next, landmarks_next, H, transformed_bumps
            item.update({'bumps': transformed_bumps})

        return item

    def get_face(self,
                 filename,
                 landmarks=None,
                 size=(cfg.CROP_SIZE, cfg.CROP_SIZE),
                 use_cache=True,
                 from_sequence=False):
        # landmarks = np.zeros((68, 2), dtype=np.float32)
        # pose = np.zeros(3, dtype=np.float32)
        crop_filepath = os.path.join(self.cropped_img_dir, filename + '.jpg')

        if use_cache and os.path.isfile(crop_filepath):
            try:
                crop = io.imread(crop_filepath)
            except OSError:
                os.remove(crop_filepath)
                return self.get_face(filename, landmarks, size, use_cache,
                                     from_sequence)
            if crop.shape[:2] != size:
                crop = cv2.resize(crop, size, interpolation=cv2.INTER_CUBIC)
            if landmarks is None:
                of_conf, landmarks, _ = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename),
                    numpy_lmFilepath=os.path.join(self.npfeature_dir,
                                                  filename))
            landmarks = face_processing.scale_landmarks_to_crop(
                landmarks, output_size=size)
        else:
            # Load image from dataset
            img_path = os.path.join(self.fullsize_img_dir, filename + '.jpg')
            img = io.imread(img_path)
            if img is None:
                raise IOError(
                    "\tError: Could not load image {}!".format(img_path))

            # load landmarks extracted with OpenFace2
            if landmarks is None:
                of_conf, landmarks, _ = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename),
                    numpy_lmFilepath=os.path.join(self.npfeature_dir,
                                                  filename),
                    from_sequence=from_sequence)
                if of_conf <= 0.0:
                    log.warning("No landmarks for image {}".format(filename))

            # crop, landmarks = face_processing.crop_bump(img, landmarks, output_size=size)
            crop, landmarks = face_processing.crop_celebHQ(img,
                                                           landmarks,
                                                           output_size=size)

            if use_cache:
                utils.io.makedirs(crop_filepath)
                io.imsave(crop_filepath, crop)

        return crop, landmarks
Ejemplo n.º 11
0
    def __init__(self,
                 root_dir=cfg.CELEBA_ROOT,
                 train=True,
                 color=True,
                 start=None,
                 max_samples=None,
                 deterministic=None,
                 crop_type='tight',
                 **kwargs):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else TEST

        self.crop_type = crop_type
        self.root_dir = root_dir
        root_dir_local = cfg.CELEBA_ROOT_LOCAL
        assert (crop_type in ['tight', 'loose', 'fullsize'])
        self.cropped_img_dir = os.path.join(root_dir_local, 'crops')
        self.fullsize_img_dir = os.path.join(root_dir, 'img_align_celeba')
        self.feature_dir = os.path.join(root_dir_local, 'features')
        self.color = color
        annotation_filename = 'list_landmarks_align_celeba.txt'

        path_annotations_mod = os.path.join(root_dir_local,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod):
            self.annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading original TXT file...')
            self.annotations = pd.read_csv(os.path.join(
                self.root_dir, 'Anno', annotation_filename),
                                           delim_whitespace=True)
            print('done.')

            # store OpenFace features in annotation dataframe
            poses = []
            confs = []
            landmarks = []
            for cnt, filename in enumerate(self.annotations.fname):
                if cnt % 1000 == 0:
                    print(cnt)
                filename_noext = os.path.splitext(filename)[0]
                conf, lms, pose = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename_noext))
                poses.append(pose)
                confs.append(conf)
                landmarks.append(lms)
            self.annotations['pose'] = poses
            self.annotations['conf'] = confs
            self.annotations['landmarks_of'] = landmarks

            # add identities to annotations
            self.identities = pd.read_csv(os.path.join(self.root_dir, 'Anno',
                                                       'identity_CelebA.txt'),
                                          delim_whitespace=True,
                                          header=None,
                                          names=['fname', 'id'])
            self.annotations = pd.merge(self.annotations,
                                        self.identities,
                                        on='fname',
                                        copy=False)

            # save annations as pickle file
            self.annotations.to_pickle(path_annotations_mod)

        # select training or test set (currently not using validation set)
        SPLIT = {
            TRAIN: (0, 162772),
            VAL: (162772, 182639),
            TEST: (182639, 202601)
        }
        self.annotations = self.annotations[
            (self.annotations.index >= SPLIT[self.mode][0])
            & (self.annotations.index < SPLIT[self.mode][1])]

        self.annotations = self.annotations.sort_values(by='id')

        print("Num. faces: {}".format(len(self.annotations)))
        if 'crops_celeba' in self.cropped_img_dir:
            min_of_conf = 0.0
        else:
            min_of_conf = 0.5
        print("Removing faces with conf < {}".format(min_of_conf))
        self.annotations = self.annotations[
            self.annotations.conf >= min_of_conf]
        print("Remaining num. faces: {}".format(len(self.annotations)))

        # max_rot_deg = 1
        # print('Limiting rotation to +-{} degrees...'.format(max_rot_deg))
        # poses = np.abs(np.stack(self.annotations.pose.values))
        # self.annotations = self.annotations[(poses[:, 0] > np.deg2rad(max_rot_deg)) |
        #                                     (poses[:, 1] > np.deg2rad(max_rot_deg)) |
        #                                     (poses[:, 2] > np.deg2rad(max_rot_deg))]
        # print(len(self.annotations))

        # limit number of samples
        st, nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st + max_samples
        self.annotations = self.annotations[st:nd]
        self._annotations = self.annotations[st:nd].copy()

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color)
Ejemplo n.º 12
0
class CelebA(td.Dataset):
    def __init__(self,
                 root_dir=cfg.CELEBA_ROOT,
                 train=True,
                 color=True,
                 start=None,
                 max_samples=None,
                 deterministic=None,
                 crop_type='tight',
                 **kwargs):

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else TEST

        self.crop_type = crop_type
        self.root_dir = root_dir
        root_dir_local = cfg.CELEBA_ROOT_LOCAL
        assert (crop_type in ['tight', 'loose', 'fullsize'])
        self.cropped_img_dir = os.path.join(root_dir_local, 'crops')
        self.fullsize_img_dir = os.path.join(root_dir, 'img_align_celeba')
        self.feature_dir = os.path.join(root_dir_local, 'features')
        self.color = color
        annotation_filename = 'list_landmarks_align_celeba.txt'

        path_annotations_mod = os.path.join(root_dir_local,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod):
            self.annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading original TXT file...')
            self.annotations = pd.read_csv(os.path.join(
                self.root_dir, 'Anno', annotation_filename),
                                           delim_whitespace=True)
            print('done.')

            # store OpenFace features in annotation dataframe
            poses = []
            confs = []
            landmarks = []
            for cnt, filename in enumerate(self.annotations.fname):
                if cnt % 1000 == 0:
                    print(cnt)
                filename_noext = os.path.splitext(filename)[0]
                conf, lms, pose = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename_noext))
                poses.append(pose)
                confs.append(conf)
                landmarks.append(lms)
            self.annotations['pose'] = poses
            self.annotations['conf'] = confs
            self.annotations['landmarks_of'] = landmarks

            # add identities to annotations
            self.identities = pd.read_csv(os.path.join(self.root_dir, 'Anno',
                                                       'identity_CelebA.txt'),
                                          delim_whitespace=True,
                                          header=None,
                                          names=['fname', 'id'])
            self.annotations = pd.merge(self.annotations,
                                        self.identities,
                                        on='fname',
                                        copy=False)

            # save annations as pickle file
            self.annotations.to_pickle(path_annotations_mod)

        # select training or test set (currently not using validation set)
        SPLIT = {
            TRAIN: (0, 162772),
            VAL: (162772, 182639),
            TEST: (182639, 202601)
        }
        self.annotations = self.annotations[
            (self.annotations.index >= SPLIT[self.mode][0])
            & (self.annotations.index < SPLIT[self.mode][1])]

        self.annotations = self.annotations.sort_values(by='id')

        print("Num. faces: {}".format(len(self.annotations)))
        if 'crops_celeba' in self.cropped_img_dir:
            min_of_conf = 0.0
        else:
            min_of_conf = 0.5
        print("Removing faces with conf < {}".format(min_of_conf))
        self.annotations = self.annotations[
            self.annotations.conf >= min_of_conf]
        print("Remaining num. faces: {}".format(len(self.annotations)))

        # max_rot_deg = 1
        # print('Limiting rotation to +-{} degrees...'.format(max_rot_deg))
        # poses = np.abs(np.stack(self.annotations.pose.values))
        # self.annotations = self.annotations[(poses[:, 0] > np.deg2rad(max_rot_deg)) |
        #                                     (poses[:, 1] > np.deg2rad(max_rot_deg)) |
        #                                     (poses[:, 2] > np.deg2rad(max_rot_deg))]
        # print(len(self.annotations))

        # limit number of samples
        st, nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st + max_samples
        self.annotations = self.annotations[st:nd]
        self._annotations = self.annotations[st:nd].copy()

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color)

    @property
    def labels(self):
        return self.annotations.id.values

    def _stats_repr(self):
        fmt_str = "    Number of identities: {}\n".format(
            self.annotations.id.nunique())
        return fmt_str

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Split: {}\n'.format(self.mode)
        fmt_str += '    Root Location: {}\n'.format(self.root_dir)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(
            tmp,
            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        fmt_str += self._stats_repr()
        return fmt_str

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

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        filename = sample.fname
        landmarks = sample.landmarks_of
        pose = sample.pose
        id = sample.id

        # crop, landmarks, pose = ds_utils.get_face(filename, self.fullsize_img_dir, self.cropped_img_dir,
        #                                           landmarks, pose, use_cache=True)
        crop, landmarks, pose, cropper = self.face_extractor.get_face(
            filename,
            self.fullsize_img_dir,
            self.cropped_img_dir,
            landmarks=landmarks,
            pose=pose,
            use_cache=True,
            detect_face=False,
            crop_type=self.crop_type,
            aligned=True)

        cropped_sample = {'image': crop, 'landmarks': landmarks, 'pose': pose}
        item = self.transform(cropped_sample)

        # face_mask = face_processing.get_face_mask(item['landmarks'], crop.shape)
        # transformed_face_mask = face_processing.CenterCrop(cfg.INPUT_SIZE)(face_mask)

        em_val_ar = np.array([[-1, 0, 0]], dtype=np.float32)

        item.update({
            'id': id,
            'fnames': filename,
            # 'face_mask': transformed_face_mask,
            'expression': em_val_ar
        })
        return item

    def get_face(self,
                 filename,
                 size=(cfg.CROP_SIZE, cfg.CROP_SIZE),
                 use_cache=True):
        print(filename)
        sample = self._annotations.loc[self._annotations.fname ==
                                       filename].iloc[0]
        landmarks = sample.landmarks_of.astype(np.float32)
        pose = sample.pose

        # if OpenFace didn't detect a face, fall back to AffectNet landmarks
        # if sample.conf < 0.1:
        #     landmarks = self.parse_landmarks(sample.facial_landmarks)

        crop, landmarks, pose = ds_utils.get_face(filename,
                                                  self.fullsize_img_dir,
                                                  self.cropped_img_dir,
                                                  landmarks,
                                                  pose,
                                                  use_cache=use_cache,
                                                  size=size)
        return crop, landmarks, pose
Ejemplo n.º 13
0
    def __init__(self, root_dir=cfg.W300_ROOT, train=True,
                 transform=None, color=True, start=None, max_samples=None,
                 deterministic=None, align_face_orientation=cfg.CROP_ALIGN_ROTATION,
                 crop_type='tight', test_split='challenging', detect_face=False, use_cache=True,
                 crop_source='bb_detector', daug=0, return_modified_images=False,
                 return_landmark_heatmaps=False, landmark_sigma=3, landmark_ids=range(68), **kwargs):

        assert(crop_type in ['fullsize', 'tight','loose'])
        test_split = test_split.lower()
        assert(test_split in ['common', 'challenging', '300w', 'full'])
        assert(crop_source in W300.CROP_SOURCES)
        lmcfg.config_landmarks('300w')

        self.start = start
        self.max_samples = max_samples
        self.use_cache = use_cache
        self.crop_source = crop_source
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.root_dir = root_dir
        self.local_root_dir = cfg.W300_ROOT_LOCAL
        self.color = color
        self.transform = transform
        self.fullsize_img_dir = os.path.join(self.root_dir, 'images')
        self.align_face_orientation = align_face_orientation
        self.detect_face = detect_face
        self.crop_type = crop_type
        self.cropped_img_dir = os.path.join(cfg.W300_ROOT_LOCAL, 'crops', crop_source)

        self.feature_dir_cnn = os.path.join(cfg.W300_ROOT_LOCAL, 'features_cnn')
        self.feature_dir_of = os.path.join(cfg.W300_ROOT_LOCAL, 'features_of')

        self.bounding_box_dir = os.path.join(cfg.W300_ROOT, 'Bounding Boxes')

        self.split = 'train' if train else test_split
        self.build_annotations(self.split)
        print("Num. images: {}".format(len(self)))

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = not train
        if self.crop_type == 'tight':
            self.transform = ds_utils.build_transform(deterministic, True, daug)
        elif self.crop_type == 'fullsize':
            self.transform = lambda x:x

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)
Ejemplo n.º 14
0
class W300(td.Dataset):

    CROP_SOURCES = ['bb_detector', 'bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface']

    def __init__(self, root_dir=cfg.W300_ROOT, train=True,
                 transform=None, color=True, start=None, max_samples=None,
                 deterministic=None, align_face_orientation=cfg.CROP_ALIGN_ROTATION,
                 crop_type='tight', test_split='challenging', detect_face=False, use_cache=True,
                 crop_source='bb_detector', daug=0, return_modified_images=False,
                 return_landmark_heatmaps=False, landmark_sigma=3, landmark_ids=range(68), **kwargs):

        assert(crop_type in ['fullsize', 'tight','loose'])
        test_split = test_split.lower()
        assert(test_split in ['common', 'challenging', '300w', 'full'])
        assert(crop_source in W300.CROP_SOURCES)
        lmcfg.config_landmarks('300w')

        self.start = start
        self.max_samples = max_samples
        self.use_cache = use_cache
        self.crop_source = crop_source
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.root_dir = root_dir
        self.local_root_dir = cfg.W300_ROOT_LOCAL
        self.color = color
        self.transform = transform
        self.fullsize_img_dir = os.path.join(self.root_dir, 'images')
        self.align_face_orientation = align_face_orientation
        self.detect_face = detect_face
        self.crop_type = crop_type
        self.cropped_img_dir = os.path.join(cfg.W300_ROOT_LOCAL, 'crops', crop_source)

        self.feature_dir_cnn = os.path.join(cfg.W300_ROOT_LOCAL, 'features_cnn')
        self.feature_dir_of = os.path.join(cfg.W300_ROOT_LOCAL, 'features_of')

        self.bounding_box_dir = os.path.join(cfg.W300_ROOT, 'Bounding Boxes')

        self.split = 'train' if train else test_split
        self.build_annotations(self.split)
        print("Num. images: {}".format(len(self)))

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = not train
        if self.crop_type == 'tight':
            self.transform = ds_utils.build_transform(deterministic, True, daug)
        elif self.crop_type == 'fullsize':
            self.transform = lambda x:x

        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)


    def build_annotations(self, split):
        import scipy.io
        import glob

        split_defs = {
            'train': [
                ('train/afw', 'afw'),
                ('train/helen', 'helen_trainset'),
                ('train/lfpw', 'lfpw_trainset')
            ],
            'common': [
                ('test/common/helen', 'helen_testset'),
                ('test/common/lfpw', 'lfpw_testset')
            ],
            'challenging': [
                ('test/challenging/ibug', 'ibug')
            ],
            'full': [
                ('test/common/helen', 'helen_testset'),
                ('test/common/lfpw', 'lfpw_testset'),
                ('test/challenging/ibug', 'ibug')
            ],
            '300w': [
                ('test/300W/01_Indoor', None),
                ('test/300W/01_Outdoor', None)
            ]
        }

        ann = []

        bboxes = []
        for id, subset in enumerate(split_defs[split]):

            im_dir, bbox_file_suffix = subset

            # get image file paths and read GT landmarks
            ext = "*.jpg"
            if 'lfpw' in im_dir or '300W' in im_dir:
                ext = "*.png"
            for img_file in sorted(glob.glob(os.path.join(self.fullsize_img_dir, im_dir, ext))):

                path_abs_noext = os.path.splitext(img_file)[0]
                filename_noext =  os.path.split(path_abs_noext)[1]
                path_rel_noext = os.path.join(im_dir, filename_noext)
                filename = os.path.split(img_file)[1]
                path_rel = os.path.join(im_dir, filename)

                # load landmarks from *.pts files
                landmarks =  ds_utils.read_300W_detection(path_abs_noext+'.pts')
                ann.append({'imgName': str(filename), 'fname': path_rel, 'landmarks': landmarks})

            # load supplied detected bounding boxes from MAT file
            if bbox_file_suffix is not None:
                subset_bboxes = scipy.io.loadmat(os.path.join(self.bounding_box_dir, 'bounding_boxes_{}.mat'.format(bbox_file_suffix)))
                for item in subset_bboxes['bounding_boxes'][0]:
                    imgName, bb_detector, bb_ground_truth = item[0][0]
                    bboxes.append({'imgName': str(imgName[0]), 'bb_detector': bb_detector[0], 'bb_ground_truth': bb_ground_truth[0]})


        self._annotations = pd.DataFrame(ann)
        if len(bboxes) > 0:
            df_bboxes = pd.DataFrame(bboxes)
            self._annotations = self._annotations.merge(df_bboxes, on='imgName', how='left')

    @property
    def labels(self):
        return None

    @property
    def annotations(self):
        return self._annotations

    @annotations.setter
    def annotations(self, new_annots):
        self._annotations = new_annots

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

    def get_bounding_box(self, sample):
        bb = sample.bb_detector if self.crop_source == 'bb_detector' else sample.bb_ground_truth

        # enlarge bounding box
        l,t,r,b = bb
        if t > b:
            t, b = b, t
        h = b-t
        assert(h >= 0)

        t_new, b_new = int(t - cfg.CROP_MOVE_TOP_FACTOR * h), int(b + cfg.CROP_MOVE_BOTTOM_FACTOR * h)
        # t_new, b_new = int(t - 0.27 * h), int(b + 0.17 * h)

        # set width of bbox same as height
        h_new = b_new - t_new
        cx = (r + l) / 2
        l_new, r_new = cx - h_new/2, cx + h_new/2

        # in case right eye is actually left of right eye...
        if l_new > r_new:
            l_new, r_new = r_new, l_new

        # extend area by crop border margins
        bbox = np.array([l_new, t_new, r_new, b_new], dtype=np.float32)

        scalef = cfg.CROP_SIZE / cfg.INPUT_SIZE
        bbox_crop = utils.geometry.scaleBB(bbox, scalef, scalef, typeBB=2)
        return bbox_crop


    def __getitem__(self, idx):
        def get_landmarks_for_crop():
            pose = np.zeros(3, dtype=np.float32)
            if self.crop_source == 'lm_openface':
                openface_filepath = os.path.join(self.feature_dir_of, os.path.splitext(filename)[0])
                est_face_center = landmarks_gt.mean(axis=0)
                of_conf, landmarks_of, pose = ds_utils.read_openface_detection(openface_filepath, expected_face_center=est_face_center)
                if of_conf < 0.01:
                    landmarks_of = landmarks_gt
                else:
                    # landmarks_of, pose = self.cropper.apply_crop_to_landmarks(landmarks_of, pose)
                    landmarks_of[:,0] -= cfg.CROP_BORDER
                    landmarks_of[:,1] -= cfg.CROP_BORDER
                landmarks  = landmarks_of
            elif self.crop_source == 'lm_cnn':
                try:
                    landmarks = np.load(os.path.join(self.feature_dir_cnn, os.path.splitext(filename)[0]+'.npy'))
                except FileNotFoundError:
                    landmarks = None
            elif self.crop_source == 'lm_ground_truth':
                landmarks = landmarks_gt
            else:
                # no landmarks -> crop using bounding boxes
                landmarks = None
            return landmarks, pose

        sample = self.annotations.iloc[idx]
        filename = sample.fname
        landmarks_gt = sample.landmarks.astype(np.float32)
        bbox = self.get_bounding_box(sample) if not self.split == '300w' else None

        landmarks_for_crop, pose = get_landmarks_for_crop()

        crop, landmarks, pose, cropper = self.face_extractor.get_face(filename, self.fullsize_img_dir, self.cropped_img_dir,
                                                                      bb=bbox, landmarks=landmarks_for_crop, pose=pose,
                                                                      use_cache=self.use_cache,
                                                                      detect_face=self.detect_face,
                                                                      crop_type=self.crop_type,
                                                                      aligned=self.align_face_orientation)

        landmarks_gt, _ = cropper.apply_to_landmarks(landmarks_gt)
        # self.show_landmarks(crop, landmarks_gt)


        cropped_sample = {'image': crop, 'landmarks': landmarks_gt, 'pose': pose}
        item = self.transform(cropped_sample)


        em_val_ar = np.array([[-1,0,0]], dtype=np.float32)

        result = self.crop_to_tensor(item)

        result.update({
            'fnames': filename,
            'expression': em_val_ar
        })

        if self.return_modified_images:
            mod_transforms = tf.Compose([fp.RandomOcclusion()])
            crop_occ = mod_transforms(item['image'])
            crop_occ = self.crop_to_tensor(crop_occ)
            result['image_mod'] = crop_occ

        # add landmark heatmaps if landmarks enabled
        if self.return_landmark_heatmaps:
            result['lm_heatmaps'] = create_landmark_heatmaps(result['landmarks'], self.landmark_sigma, self.landmark_ids)

        return result

    def show_landmarks(self, img, landmarks):
        for lm in landmarks:
            lm_x, lm_y = lm[0], lm[1]
            cv2.circle(img, (int(lm_x), int(lm_y)), 3, (0, 0, 255), -1)
        cv2.imshow('landmarks', cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        cv2.waitKey(0)
Ejemplo n.º 15
0
    def __init__(self,
                 root_dir=cfg.VGGFACE2_ROOT,
                 train=True,
                 color=True,
                 start=None,
                 max_samples=None,
                 deterministic=None,
                 min_conf=cfg.MIN_OPENFACE_CONFIDENCE,
                 use_cache=True,
                 crop_source='bb_ground_truth',
                 detect_face=False,
                 align_face_orientation=True,
                 return_landmark_heatmaps=False,
                 return_modified_images=False,
                 daug=0,
                 landmark_sigma=None,
                 landmark_ids=None,
                 **kwargs):

        assert (crop_source in [
            'bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface'
        ])

        self.mode = TRAIN if train else VAL

        self.face_extractor = FaceExtractor()
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.color = color
        self.crop_source = crop_source
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.root_dir = root_dir
        root_dir_local = cfg.VGGFACE2_ROOT_LOCAL
        split_subfolder = 'train' if train else 'test'
        self.cropped_img_dir = os.path.join(root_dir_local, split_subfolder,
                                            'crops', crop_source)
        self.fullsize_img_dir = os.path.join(root_dir, split_subfolder, 'imgs')
        self.feature_dir = os.path.join(root_dir_local, split_subfolder,
                                        'features')
        annotation_filename = 'loose_bb_{}.csv'.format(split_subfolder)
        # annotation_filename = 'loose_landmark_{}.csv'.format(split_subfolder)

        # self.path_annotations_mod = os.path.join(root_dir_local, annotation_filename + '.mod_full_of.pkl')
        self.path_annotations_mod = os.path.join(
            root_dir_local, annotation_filename + '.mod_full.pkl')
        if os.path.isfile(self.path_annotations_mod):
            print('Reading pickle file...')
            self.annotations = pd.read_pickle(self.path_annotations_mod)
            print('done.')
        else:
            print('Reading CSV file...')
            self.annotations = pd.read_csv(
                os.path.join(self.root_dir, 'bb_landmark',
                             annotation_filename))
            print('done.')

            of_confs, poses, landmarks = [], [], []
            self.annotations = self.annotations[0:4000000]
            self.annotations = self.annotations[self.annotations.H > 80]
            print("Number of images: {}".format(len(self)))

            def get_face_height(lms):
                return lms[8, 1] - lms[27, 1]

            read_openface_landmarks = True
            if read_openface_landmarks:
                for cnt, filename in enumerate(self.annotations.NAME_ID):
                    filename_noext = os.path.splitext(filename)[0]

                    bb = self.annotations.iloc[cnt][1:5].values
                    expected_face_center = [
                        bb[0] + bb[2] / 2.0, bb[1] + bb[3] / 2.0
                    ]

                    conf, lms, pose, num_faces = ds_utils.read_openface_detection(
                        os.path.join(self.feature_dir, filename_noext),
                        expected_face_center=expected_face_center,
                        use_cache=True,
                        return_num_faces=True)

                    if num_faces > 1:
                        print("Deleting extracted crop for {}...".format(
                            filename))
                        cache_filepath = os.path.join(self.cropped_img_dir,
                                                      'tight',
                                                      filename + '.jpg')
                        if os.path.isfile(cache_filepath):
                            os.remove(cache_filepath)

                        # numpy_lmfile = os.path.join(self.feature_dir, filename) + '.npz'
                        # if os.path.isfile(numpy_lmfile):
                        #     os.remove(numpy_lmfile)

                    of_confs.append(conf)
                    landmarks.append(lms)
                    poses.append(pose)
                    if (cnt + 1) % 10000 == 0:
                        log.info(cnt + 1)
                    # if (cnt+1) % 1000 == 0:
                    #     print('saving annotations...')
                    #     self.annotations.to_pickle(self.path_annotations_mod)
                self.annotations['pose'] = poses
                self.annotations['of_conf'] = of_confs
                self.annotations['landmarks_of'] = landmarks
                # self.annotations['face_height'] = self.annotations.landmarks_of.map(get_face_height)

            # assign new continuous ids to persons (0, range(n))
            print("Creating id labels...")
            _ids = self.annotations.NAME_ID
            _ids = _ids.map(lambda x: int(x.split('/')[0][1:]))
            self.annotations['ID'] = _ids

            # unique_ids = _ids.unique()
            # uid2idx = dict(zip(unique_ids, range(1,len(unique_ids)+1)))
            # self.annotations['ID'] = _ids.map(uid2idx)

            self.annotations.to_pickle(self.path_annotations_mod)

        select_subset = False
        if select_subset:
            print("Number of images: {}".format(len(self)))
            self.annotations = self.annotations[
                self.annotations.of_conf > min_conf]
            print("Number of images: {}".format(len(self)))
            min_rot_deg = 0
            max_rot_deg = 90
            print('Limiting rotation to +-[{}-{}] degrees...'.format(
                min_rot_deg, max_rot_deg))
            poses = np.abs(np.stack(self.annotations.pose.values))
            self.annotations = self.annotations[
                (poses[:, 0] < np.deg2rad(max_rot_deg))
                & (poses[:, 1] < np.deg2rad(max_rot_deg)) &
                (poses[:, 2] < np.deg2rad(max_rot_deg))]
            # self.annotations = self.annotations[(np.deg2rad(min_rot_deg) < poses[:, 0]) |
            #                                     (np.deg2rad(min_rot_deg) < poses[:, 1])]

        min_face_height = 100
        print(
            'Removing faces with height <={:.2f}px...'.format(min_face_height))
        self.annotations = self.annotations[
            self.annotations.H > min_face_height]
        print("Number of images: {}".format(len(self)))

        # width = self.annotations.W
        # height = self.annotations.H
        # ratio = width / height
        # self.annotations = self.annotations[(ratio > 0.5) & (ratio < 0.60)]
        # self.annotations = self.annotations[ratio < 0.65]
        # self.annotations = self.annotations[ratio > 0.9]

        # FIXME: shuffle for find_similar_images
        # self.annotations = self.annotations[:1000000]
        # from sklearn.utils import shuffle
        # self.annotations = shuffle(self.annotations, random_state=2)
        #############3

        # limit number of samples
        st, nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st + max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color,
                                                  daug)

        print("Number of images: {}".format(len(self)))
        print("Number of identities: {}".format(self.annotations.ID.nunique()))
Ejemplo n.º 16
0
    def __init__(self, root_dir=cfg.VGGFACE2_ROOT, train=True, color=True, start=None,
                 max_samples=None, deterministic=None, min_conf=cfg.MIN_OPENFACE_CONFIDENCE, use_cache=True,
                 crop_source='bb_ground_truth', detect_face=False, align_face_orientation=True,
                 return_landmark_heatmaps=False, return_modified_images=False,
                 daug=0, landmark_sigma=None, landmark_ids=None, **kwargs):

        assert(crop_source in ['bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface'])

        self.mode = TRAIN if train else VAL

        self.face_extractor = FaceExtractor()
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.color = color
        self.crop_source = crop_source
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.root_dir = root_dir
        root_dir_local = cfg.VGGFACE2_ROOT_LOCAL
        split_subfolder = 'train' if train else 'test'
        crop_folder = 'crops'
        if cfg.INPUT_SIZE == 128:
            crop_folder += '_128'
        self.cropped_img_dir = os.path.join(root_dir_local, split_subfolder, crop_folder, crop_source)
        self.fullsize_img_dir = os.path.join(root_dir, split_subfolder, 'imgs')
        self.feature_dir = os.path.join(root_dir_local, split_subfolder, 'features')
        annotation_filename = 'loose_bb_{}.csv'.format(split_subfolder)
        # annotation_filename = 'loose_landmark_{}.csv'.format(split_subfolder)

        # self.path_annotations_mod = os.path.join(root_dir_local, annotation_filename + '.mod_full_of.pkl')
        self.path_annotations_mod = os.path.join(root_dir_local, annotation_filename + '.mod_full.pkl')
        if os.path.isfile(self.path_annotations_mod):
            print('Reading pickle file...')
            self.annotations = pd.read_pickle(self.path_annotations_mod)
            print('done.')
        else:
            print('Reading CSV file...')
            self.annotations = pd.read_csv(os.path.join(self.root_dir, 'bb_landmark', annotation_filename))
            print('done.')

            of_confs, poses, landmarks = [], [], []
            self.annotations = self.annotations[0:4000000]
            self.annotations = self.annotations[self.annotations.H > 80]
            print("Number of images: {}".format(len(self)))

            def get_face_height(lms):
                return lms[8,1] - lms[27,1]

            read_openface_landmarks = True
            if read_openface_landmarks:
                for cnt, filename in enumerate(self.annotations.NAME_ID):
                    filename_noext = os.path.splitext(filename)[0]

                    bb = self.annotations.iloc[cnt][1:5].values
                    expected_face_center = [bb[0] + bb[2] / 2.0, bb[1] + bb[3] / 2.0]

                    conf, lms, pose, num_faces  = ds_utils.read_openface_detection(os.path.join(self.feature_dir, filename_noext),
                                                                       expected_face_center=expected_face_center,
                                                                       use_cache=True, return_num_faces=True)

                    if num_faces > 1:
                        print("Deleting extracted crop for {}...".format(filename))
                        cache_filepath = os.path.join(self.cropped_img_dir, 'tight', filename + '.jpg')
                        if os.path.isfile(cache_filepath):
                            os.remove(cache_filepath)

                    of_confs.append(conf)
                    landmarks.append(lms)
                    poses.append(pose)
                    if (cnt+1) % 10000 == 0:
                        log.info(cnt+1)
                self.annotations['pose'] = poses
                self.annotations['of_conf'] = of_confs
                self.annotations['landmarks_of'] = landmarks

            # assign new continuous ids to persons (0, range(n))
            print("Creating id labels...")
            _ids = self.annotations.NAME_ID
            _ids = _ids.map(lambda x: int(x.split('/')[0][1:]))
            self.annotations['ID'] = _ids

            self.annotations.to_pickle(self.path_annotations_mod)

        min_face_height = 100
        print('Removing faces with height <={:.2f}px...'.format(min_face_height))
        self.annotations = self.annotations[self.annotations.H > min_face_height]
        print("Number of images: {}".format(len(self)))

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color, daug)

        print("Number of images: {}".format(len(self)))
        print("Number of identities: {}".format(self.annotations.ID.nunique()))
Ejemplo n.º 17
0
class AffectNet(td.Dataset):

    classes = CLASS_NAMES
    colors = [
        'tab:gray', 'tab:orange', 'tab:brown', 'tab:pink', 'tab:cyan',
        'tab:olive', 'tab:red', 'tab:blue'
    ]
    markers = ['s', 'o', '>', '<', '^', 'v', 'P', 'd']

    def __init__(self,
                 root_dir=cfg.AFFECTNET_ROOT,
                 train=True,
                 transform=None,
                 crop_type='tight',
                 color=True,
                 start=None,
                 max_samples=None,
                 outlier_threshold=None,
                 deterministic=None,
                 use_cache=True,
                 detect_face=False,
                 align_face_orientation=False,
                 min_conf=cfg.MIN_OPENFACE_CONFIDENCE,
                 daug=0,
                 return_landmark_heatmaps=False,
                 landmark_sigma=9,
                 landmark_ids=range(68),
                 return_modified_images=False,
                 crop_source='lm_openface',
                 **kwargs):
        assert (crop_type in ['fullsize', 'tight', 'loose'])
        assert (crop_source in [
            'bb_ground_truth', 'lm_ground_truth', 'lm_cnn', 'lm_openface'
        ])

        self.face_extractor = FaceExtractor()

        self.mode = TRAIN if train else VAL

        self.crop_source = crop_source
        self.use_cache = use_cache
        self.detect_face = detect_face
        self.align_face_orientation = align_face_orientation
        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.start = start
        self.max_samples = max_samples

        self.root_dir = root_dir
        self.crop_type = crop_type
        self.color = color
        self.outlier_threshold = outlier_threshold
        self.transform = transform
        self.fullsize_img_dir = os.path.join(self.root_dir,
                                             'cropped_Annotated')
        self.cropped_img_dir = os.path.join(self.root_dir, 'crops',
                                            crop_source)
        self.feature_dir = os.path.join(self.root_dir, 'features')

        annotation_filename = 'training' if train else 'validation'
        path_annotations_mod = os.path.join(root_dir,
                                            annotation_filename + '.mod.pkl')
        if os.path.isfile(path_annotations_mod):
            print('Reading pickle file...')
            self._annotations = pd.read_pickle(path_annotations_mod)
        else:
            print('Reading CSV file...')
            self._annotations = pd.read_csv(
                os.path.join(root_dir, annotation_filename + '.csv'))
            print('done.')

            # drop non-faces
            self._annotations = self._annotations[
                self._annotations.expression < 8]

            # Samples in annotation file are somewhat clustered by expression.
            # Shuffle to create a more even distribution.
            # NOTE: deterministic, always creates the same order
            if train:
                from sklearn.utils import shuffle
                self._annotations = shuffle(self._annotations, random_state=2)

                # remove samples with inconsistent expression<->valence/arousal values
                self._remove_outliers()

            poses = []
            confs = []
            landmarks = []
            for cnt, filename in enumerate(
                    self._annotations.subDirectory_filePath):
                if cnt % 1000 == 0:
                    print(cnt)
                filename_noext = os.path.splitext(filename)[0]
                conf, lms, pose = ds_utils.read_openface_detection(
                    os.path.join(self.feature_dir, filename_noext))
                poses.append(pose)
                confs.append(conf)
                landmarks.append(lms)
            self._annotations['pose'] = poses
            self._annotations['conf'] = confs
            self._annotations['landmarks_of'] = landmarks
            # self.annotations.to_csv(path_annotations_mod, index=False)
            self._annotations.to_pickle(path_annotations_mod)

        poses = np.abs(np.stack(self._annotations.pose.values))

        only_good_image_for_training = True
        if train and only_good_image_for_training:
            print(len(self._annotations))

            min_rot_deg = 30
            max_rot_deg = 90
            # print('Limiting rotation to +-[{}-{}] degrees...'.format(min_rot_deg, max_rot_deg))
            # self._annotations = self._annotations[(poses[:, 0] < np.deg2rad(max_rot_deg)) &
            #                                       (poses[:, 1] < np.deg2rad(max_rot_deg)) &
            #                                       (poses[:, 2] < np.deg2rad(max_rot_deg))]
            # self._annotations = self._annotations[(np.deg2rad(min_rot_deg) < poses[:, 0]) |
            #                                       (np.deg2rad(min_rot_deg) < poses[:, 1])]
            # self._annotations = self._annotations[np.deg2rad(min_rot_deg) < poses[:, 1] ]

            print(len(self._annotations))

            # print('Removing OpenFace confs <={:.2f}...'.format(min_conf))
            # self._annotations = self._annotations[self._annotations.conf > cfg.MIN_OPENFACE_CONFIDENCE]
            # print(len(self._annotations))

            # select by Valence/Arousal
            # min_arousal = 0.0
            # print('Removing arousal <={:.2f}...'.format(min_arousal))
            # self._annotations = self._annotations[self._annotations.arousal > min_arousal]
            # print(len(self._annotations))

        # There is (at least) one missing image in the dataset. Remove by checking face width:
        self._annotations = self._annotations[self._annotations.face_width > 0]

        # self._annotations_balanced = self._annotations
        # self.filter_labels(label_dict_exclude={'expression': 0})
        # self.filter_labels(label_dict_exclude={'expression': 1})
        # self._annotations = self._annotations[self._annotations.arousal > 0.2]

        self.rebalance_classes()

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, self.color,
                                                  daug)

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor()]
        transforms += [fp.Normalize([0.518, 0.418, 0.361],
                                    [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)

    def filter_labels(self, label_dict=None, label_dict_exclude=None):
        if label_dict is not None:
            print("Applying include filter to labels: {}".format(label_dict))
            for k, v in label_dict.items():
                self.annotations = self.annotations[self.annotations[k] == v]
        if label_dict_exclude is not None:
            print("Applying exclude filter to labels: {}".format(
                label_dict_exclude))
            for k, v in label_dict_exclude.items():
                self.annotations = self.annotations[self.annotations[k] != v]
        print("  Number of images: {}".format(len(self.annotations)))

    def rebalance_classes(self,
                          max_images_per_class=MAX_IMAGES_PER_EXPRESSION):
        if self.mode == TRAIN:
            # balance class sized if neccessary
            print('Limiting number of images to {} per class...'.format(
                max_images_per_class))
            # self._annotations = self._annotations.groupby('expression').head(5000)
            from sklearn.utils import shuffle
            self._annotations['cls_idx'] = self._annotations.groupby(
                'expression').cumcount()
            self._annotations = shuffle(self._annotations)
            self._annotations_balanced = self._annotations[
                self._annotations.cls_idx < max_images_per_class]
            print(len(self._annotations_balanced))
        else:
            self._annotations_balanced = self._annotations

        # limit number of samples
        st, nd = 0, None
        if self.start is not None:
            st = self.start
        if self.max_samples is not None:
            nd = st + self.max_samples
        self._annotations_balanced = self._annotations_balanced[st:nd]

    @property
    def labels(self):
        return self.annotations['expression'].values

    @property
    def heights(self):
        return self.annotations.face_height.values

    @property
    def widths(self):
        return self.annotations.face_width.values

    @property
    def annotations(self):
        return self._annotations_balanced

    @annotations.setter
    def annotations(self, new_annots):
        self._annotations_balanced = new_annots

    def print_stats(self):
        print(self._stats_repr())

    def _stats_repr(self):
        labels = self.annotations.expression
        fmt_str = "    Class sizes:\n"
        for id in np.unique(labels):
            count = len(np.where(labels == id)[0])
            fmt_str += "      {:<6} ({:.2f}%)\t({})\n".format(
                count, 100.0 * count / self.__len__(), self.classes[id])
        fmt_str += "    --------------------------------\n"
        fmt_str += "      {:<6}\n".format(len(labels))
        return fmt_str

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Split: {}\n'.format(self.mode)
        # fmt_str += '    Root Location: {}\n'.format(self.root_dir)
        # tmp = '    Transforms (if any): '
        # fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        fmt_str += self._stats_repr()
        return fmt_str

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

    def get_class_sizes(self):
        groups = self.annotations.groupby(by='expression')
        return groups.size().values

    def _remove_outliers(self):
        if self.outlier_threshold is None:
            return
        from utils.exprec import calc_mahalanobis_covs, get_expression_dists
        covs = calc_mahalanobis_covs(self.annotations)

        VA = self._annotations.as_matrix(columns=['valence', 'arousal'])
        true_class = self._annotations['expression'].values
        dists = get_expression_dists(VA, covs)
        class_preds = np.argmin(dists, axis=1)
        true_class_dist = dists[range(len(dists)), tuple(true_class)]
        self._annotations['dist'] = true_class_dist
        self._annotations['class_pred'] = class_preds
        count_before = len(self._annotations)
        self._annotations = self._annotations.loc[
            self._annotations['dist'] < self.outlier_threshold]
        print("Removed {} outliers from dataset (th={}).".format(
            count_before - len(self._annotations), self.outlier_threshold))

    def parse_landmarks(self, landmarks):
        try:
            vals = [float(s) for s in landmarks.split(';')]
            return np.array([(x, y) for x, y in zip(vals[::2], vals[1::2])],
                            dtype=np.float32)
        except:
            raise ValueError("Invalid landmarks {}".format(landmarks))

    def get_bounding_box(self, sample):
        l, t, w, h = sample.face_x, sample.face_y, sample.face_width, sample.face_height
        r, b = l + w, t + h

        # return np.array([l,t,r,b], dtype=np.float32)

        # enlarge bounding box
        if t > b:
            t, b = b, t
        h = b - t
        assert (h >= 0)
        t_new, b_new = int(t + 0.05 * h), int(b + 0.25 * h)

        # set width of bbox same as height
        h_new = b_new - t_new
        cx = (r + l) / 2
        l_new, r_new = cx - h_new / 2, cx + h_new / 2
        # in case right eye is actually left of right eye...
        if l_new > r_new:
            l_new, r_new = r_new, l_new

        # extend area by crop border margins
        bbox = np.array([l_new, t_new, r_new, b_new], dtype=np.float32)
        scalef = cfg.CROP_SIZE / cfg.INPUT_SIZE
        bbox_crop = utils.geometry.scaleBB(bbox, scalef, scalef, typeBB=2)

        return bbox_crop

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        filename = sample.subDirectory_filePath
        pose = sample.pose
        bb = None
        landmarks_for_crop = None
        landmarks_to_return = self.parse_landmarks(sample.facial_landmarks)

        if self.crop_source == 'bb_ground_truth':
            bb = self.get_bounding_box(sample)
        elif self.crop_source == 'lm_ground_truth':
            landmarks_for_crop = landmarks_to_return
        elif self.crop_source == 'lm_openface':
            of_conf, landmarks_for_crop = sample.conf, sample.landmarks_of
            # if OpenFace didn't detect a face, fall back to AffectNet landmarks
            if sample.conf <= 0.1:
                try:
                    landmarks_for_crop = self.parse_landmarks(
                        sample.facial_landmarks)
                except ValueError:
                    pass

        try:
            crop, landmarks, pose, cropper = self.face_extractor.get_face(
                filename,
                self.fullsize_img_dir,
                self.cropped_img_dir,
                landmarks=landmarks_for_crop,
                bb=bb,
                pose=pose,
                use_cache=self.use_cache,
                detect_face=False,
                crop_type=self.crop_type,
                aligned=self.align_face_orientation)
        except AssertionError:
            print(filename)
            raise

        landmarks, _ = cropper.apply_to_landmarks(landmarks_to_return)
        # vis.show_landmarks(crop, landmarks, title='lms affectnet', wait=0, color=(0,0,255))

        cropped_sample = {'image': crop, 'landmarks': landmarks, 'pose': pose}
        item = self.transform(cropped_sample)

        em_val_ar = np.array(
            [[sample.expression, sample.valence, sample.arousal]],
            dtype=np.float32)

        result = self.crop_to_tensor(item)

        result.update({'id': 0, 'fnames': filename, 'expression': em_val_ar})

        if self.return_modified_images:
            mod_transforms = tf.Compose([fp.RandomOcclusion()])
            crop_occ = mod_transforms(item['image'])
            crop_occ = self.crop_to_tensor(crop_occ)
            result['image_mod'] = crop_occ

        if self.return_landmark_heatmaps:
            result['lm_heatmaps'] = create_landmark_heatmaps(
                result['landmarks'], self.landmark_sigma, self.landmark_ids)

        return result

    def get_face(self,
                 filename,
                 size=(cfg.CROP_SIZE, cfg.CROP_SIZE),
                 use_cache=True):
        sample = self._annotations.loc[self._annotations.subDirectory_filePath
                                       == filename].iloc[0]
        landmarks = sample.landmarks_of.astype(np.float32)
        pose = sample.pose

        # if OpenFace didn't detect a face, fall back to AffectNet landmarks
        if sample.conf <= 0.9:
            landmarks = self.parse_landmarks(sample.facial_landmarks)

        crop, landmarks, pose, _ = self.face_extractor.get_face(
            filename,
            self.fullsize_img_dir,
            self.cropped_img_dir,
            crop_type='tight',
            landmarks=landmarks,
            pose=pose,
            use_cache=True,
            detect_face=False,
            size=size)

        return crop, landmarks, pose

    def show_landmarks(self, img, landmarks):
        for lm in landmarks:
            lm_x, lm_y = lm[0], lm[1]
            cv2.circle(img, (int(lm_x), int(lm_y)), 3, (0, 0, 255), -1)
        cv2.imshow('landmarks', cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        cv2.waitKey(0)
Ejemplo n.º 18
0
class AFLW(td.Dataset):

    def __init__(self, root_dir=cfg.AFLW_ROOT, train=True, color=True, start=None,
                 max_samples=None, deterministic=None, use_cache=True,
                 daug=0, return_modified_images=False, test_split='full', align_face_orientation=True,
                 return_landmark_heatmaps=False, landmark_sigma=9, landmark_ids=range(19), **kwargs):

        assert test_split in ['full', 'frontal']
        from utils.face_extractor import FaceExtractor
        self.face_extractor = FaceExtractor()

        self.use_cache = use_cache
        self.align_face_orientation = align_face_orientation

        self.return_landmark_heatmaps = return_landmark_heatmaps
        self.return_modified_images = return_modified_images
        self.landmark_sigma = landmark_sigma
        self.landmark_ids = landmark_ids

        self.mode = TRAIN if train else VAL

        self.root_dir = root_dir
        root_dir_local = cfg.AFLW_ROOT_LOCAL
        self.fullsize_img_dir = os.path.join(root_dir, 'data/flickr')
        self.cropped_img_dir = os.path.join(root_dir_local, 'crops')
        self.feature_dir = os.path.join(root_dir_local,  'features')
        self.color = color

        annotation_filename = os.path.join(cfg.AFLW_ROOT_LOCAL, 'alfw.pkl')
        self.annotations_original = pd.read_pickle(annotation_filename)
        print("Number of images: {}".format(len(self.annotations_original)))

        self.frontal_only = test_split == 'frontal'
        self.make_split(train, self.frontal_only)

        # limit number of samples
        st,nd = 0, None
        if start is not None:
            st = start
        if max_samples is not None:
            nd = st+max_samples
        self.annotations = self.annotations[st:nd]

        if deterministic is None:
            deterministic = self.mode != TRAIN
        self.transform = ds_utils.build_transform(deterministic, True, daug)

        transforms = [fp.CenterCrop(cfg.INPUT_SIZE)]
        transforms += [fp.ToTensor() ]
        transforms += [fp.Normalize([0.518, 0.418, 0.361], [1, 1, 1])]  # VGGFace(2)
        self.crop_to_tensor = tf.Compose(transforms)

        print("Number of images: {}".format(len(self)))
        # print("Number of identities: {}".format(self.annotations.id.nunique()))

    @property
    def labels(self):
        return self.annotations.ID.values

    @property
    def heights(self):
        return self.annotations.face_h.values

    @property
    def widths(self):
        return self.annotations.face_w.values

    def make_split(self, train, only_frontal):
        import scipy.io
        # Additional annotations from http://mmlab.ie.cuhk.edu.hk/projects/compositional.html
        annots = scipy.io.loadmat(os.path.join(cfg.AFLW_ROOT_LOCAL, 'AFLWinfo_release.mat'))

        train_ids, test_ids = annots['ra'][0][:20000] - 1, annots['ra'][0][20000:] - 1
        ids = annots['ra'][0] - 1


        # merge original and additional annotations

        lms = annots['data'][ids]
        lms = np.dstack((lms[:,:19], lms[:, 19:]))
        lms_list = [l for l in lms]
        mask_new = annots['mask_new'][ids]

        # mask_all_lms_visible = np.stack(mask_new).min(axis=1) == 1

        bbox = annots['bbox'][ids]
        x1, x2, y1, y2 = bbox[:,0], bbox[:,1], bbox[:, 2], bbox[:, 3]
        fnames = [f[0][0] for f in annots['nameList'][ids]]
        annotations_additional = pd.DataFrame({
                                               'fname':fnames,
                                               'ra': ids,
                                               'landmarks_full':lms_list,
                                               'masks': [m for m in mask_new],
                                               'face_x': x1,
                                               'face_y': y1,
                                               'face_w': x2 - x1,
                                               'face_h': y2 - y1
            })

        ad = annotations_additional
        ao = self.annotations_original

        # self.annotations_test = self.annotations_original[self.annotations.fname.isin(fnames)]
        pd.set_option('display.expand_frame_repr', False)
        self.annotations = pd.merge(ad, ao, on=['fname',
                                                'face_x',
                                                'face_y',
                                                'face_w',
                                                'face_h'
                                                ])
        self.annotations = self.annotations.sort_values('ra')


        split_ids = train_ids if train else test_ids
        self.annotations = self.annotations[self.annotations.ra.isin(split_ids)]

        if not train and only_frontal:
            mask_all_lms_visible = np.stack(self.annotations.masks.values).min(axis=1) == 1
            self.annotations = self.annotations[mask_all_lms_visible]
            print(len(self.annotations))


    def get_bounding_box(self, sample):
        l,t,w,h = sample.face_x, sample.face_y, sample.face_w, sample.face_h
        r, b = l + w, t + h

        # enlarge bounding box
        if t > b:
            t, b = b, t
        h = b-t
        assert(h >= 0)
        t_new, b_new = int(t - 0.05 * h), int(b + 0.08 * h)

        # set width of bbox same as height
        h_new = b_new - t_new
        cx = (r + l) / 2
        l_new, r_new = cx - h_new/2, cx + h_new/2
        # in case right eye is actually left of right eye...
        if l_new > r_new:
            l_new, r_new = r_new, l_new

        # extend area by crop border margins
        bbox = np.array([l_new, t_new, r_new, b_new], dtype=np.float32)
        scalef = cfg.CROP_SIZE / cfg.INPUT_SIZE
        bbox_crop = utils.geometry.scaleBB(bbox, scalef, scalef, typeBB=2)
        return bbox_crop

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

    def __getitem__(self, idx):
        sample = self.annotations.iloc[idx]
        # assert sample.fname == sample.fname_full

        face_id = sample.ra
        filename  = sample.fname
        bb = self.get_bounding_box(sample)

        try:
            crop, landmarks, pose, cropper = self.face_extractor.get_face(filename,
                                                                          self.fullsize_img_dir,
                                                                          self.cropped_img_dir,
                                                                          bb=bb,
                                                                          use_cache=self.use_cache,
                                                                          id=face_id)
        except:
            print(filename)
            raise

        landmarks, _ = cropper.apply_to_landmarks(sample.landmarks_full)
        # vis.show_landmarks(crop, landmarks, title='lms aflw', wait=0, color=(0,0,255))

        cropped_sample = {'image': crop,
                          'landmarks': landmarks.astype(np.float32),
                          # 'landmarks': np.zeros((68,2), dtype=np.float32),
                          'pose':  np.zeros(3, dtype=np.float32)}
        item = self.transform(cropped_sample)

        em_val_ar = np.array([[-1,0,0]], dtype=np.float32)

        result = self.crop_to_tensor(item)

        result.update({
            'fnames': filename,
            'expression': em_val_ar,
            'id': 0
        })

        if self.return_modified_images:
            mod_transforms = tf.Compose([fp.RandomOcclusion()])
            crop_occ = mod_transforms(item['image'])
            crop_occ = self.crop_to_tensor(crop_occ)
            result['image_mod'] = crop_occ

        if self.return_landmark_heatmaps:
            result['lm_heatmaps'] = create_landmark_heatmaps(result['landmarks'], self.landmark_sigma, self.landmark_ids)

        return result