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
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
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
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)
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
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
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)
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