def __getitem__(self, idx):
"""Returns sample (image, index)"""
img = None
try:
img = cv.imread(self.samples_info[idx]['path'], cv.IMREAD_COLOR)
bbox = self.samples_info[idx]['bbox']
landmarks = self.samples_info[idx]['landmarks']
if bbox is not None or landmarks is not None:
if landmarks is not None:
landmarks = np.array(landmarks).reshape(5, -1)
landmarks[:,0] = landmarks[:,0]*bbox[2] + bbox[0]
landmarks[:,1] = landmarks[:,1]*bbox[3] + bbox[1]
img = FivePointsAligner.align(img, landmarks.reshape(-1), d_size=(bbox[2], bbox[3]),
normalized=False, show=False)
if bbox is not None and landmarks is None:
img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]
except BaseException:
print('Corrupted image!', self.samples_info[idx])
img = np.zeros((128, 128, 3), dtype='uint8')
if self.transform:
img = self.transform(img)
return {'img': img, 'idx': idx}
def __getitem__(self, idx):
"""Returns sample (image, class id, image id) by index"""
img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)
if self.landmarks_training:
landmarks = self.samples_info[idx][-1]
bbox = self.samples_info[idx][-2]
img = img[bbox[1]:bbox[1] + bbox[3], bbox[0]:bbox[0] + bbox[2]]
landmarks = [(float(landmarks[2 * i] - bbox[0]) / bbox[2],
float(landmarks[2 * i + 1] - bbox[1]) / bbox[3])
for i in range(len(landmarks) // 2)]
data = {'img': img, 'landmarks': np.array(landmarks)}
if self.transform:
data = self.transform(data)
return data
if self.have_landmarks:
landmarks = self.samples_info[idx][-1]
img = FivePointsAligner.align(img,
landmarks,
d_size=(200, 200),
normalized=False)
if self.transform:
img = self.transform(img)
return {
'img': img,
'label': self.samples_info[idx][1],
'instance': self.samples_info[idx][2]
}
def show_item(self, index):
"""Saves a pair with a given index to disk"""
path_1, path_2, _, _, _ = self.pairs[index]
img1 = cv.imread(path_1)
img2 = cv.imread(path_2)
if self.use_landmarks:
landmarks1 = np.array(self.landmarks[path_1[path_1.rfind('/') +
1:]]).reshape(-1)
landmarks2 = np.array(self.landmarks[path_2[path_2.rfind('/') +
1:]]).reshape(-1)
img1 = FivePointsAligner.align(img1, landmarks1)
img2 = FivePointsAligner.align(img2, landmarks2)
else:
img1 = cv.resize(img1, (400, 400))
img2 = cv.resize(img2, (400, 400))
cv.imwrite('misclassified_{}.jpg'.format(index),
np.hstack([img1, img2]))
def __getitem__(self, idx):
img = cv.imread(self.samples_info[idx][0])
if self.use_landmarks:
img = FivePointsAligner.align(img, self.samples_info[idx][2],
d_size=(200, 200), normalized=True, show=False)
if self.transform:
img = self.transform(img)
return {'img': img, 'label': int(self.samples_info[idx][1])}
def compute_embeddings_lfw_ie(args,
dataset,
model,
batch_size=1,
dump_embeddings=False,
pdist=cosine,
flipped_embeddings=False,
lm_model=None):
"""Computes embeddings of all images from the LFW dataset using Inference Engine"""
assert batch_size == 1
scores_with_gt = []
for batch_idx, data in enumerate(tqdm(dataset, 'Computing embeddings')):
images_1 = data['img1']
images_2 = data['img2']
if lm_model:
lm_input_size = tuple(lm_model.get_input_shape()[2:])
landmarks_1 = lm_model.forward(cv.resize(
images_1, lm_input_size)).reshape(-1)
images_1 = FivePointsAligner.align(images_1,
landmarks_1,
*images_1.shape[:2],
normalize=False,
show=False)
landmarks_2 = lm_model.forward(cv.resize(
images_2, lm_input_size)).reshape(-1)
images_2 = FivePointsAligner.align(images_2,
landmarks_2,
*images_2.shape[:2],
normalize=False)
is_same = data['is_same']
emb_1 = model.forward(images_1).reshape(-1)
emb_2 = model.forward(images_2).reshape(-1)
score = pdist(emb_1, emb_2)
scores_with_gt.append({
'score': score,
'is_same': is_same,
'idx': batch_idx * batch_size
})
return scores_with_gt
def _load_img(self, img_path):
"""Loads an image from dist, then performs face alignment and applies transform"""
img = cv.imread(img_path, cv.IMREAD_COLOR)
if self.use_landmarks:
landmarks = np.array(self.landmarks[img_path[img_path.rfind('/') +
1:]]).reshape(-1)
img = FivePointsAligner.align(img, landmarks, show=False)
if self.transform is None:
return img
return self.transform(img)
def __getitem__(self, idx):
"""Returns sample (image, class id, image id) by index"""
img = cv.imread(self.samples_info[idx][0], cv.IMREAD_COLOR)
landmarks = self.samples_info[idx][-1]
img = FivePointsAligner.align(img,
landmarks,
d_size=(200, 200),
normalized=True,
show=False)
if self.transform:
img = self.transform(img)
return {
'img': img,
'label': self.samples_info[idx][1],
'instance': self.samples_info[idx][2]
}
def test_align_image(self):
"""Synthetic test for alignment function"""
image = np.zeros((128, 128, 3), dtype=np.float32)
for point in FivePointsAligner.ref_landmarks:
point_scaled = point * [128, 128]
cv.circle(image, tuple(point_scaled.astype(np.int)), 5,
(255, 255, 255), cv.FILLED)
transform = RandomRotate(40., p=1.)
rotated_data = transform({
'img': image,
'landmarks': FivePointsAligner.ref_landmarks
})
aligned_image = FivePointsAligner.align(rotated_data['img'], \
rotated_data['landmarks'].reshape(-1),
d_size=(128, 128), normalized=True)
for point in FivePointsAligner.ref_landmarks:
point_scaled = (point * [128, 128]).astype(np.int)
check_sum = np.mean(
aligned_image[point_scaled[1] - 3:point_scaled[1] + 3,
point_scaled[0] - 3:point_scaled[0] + 3])
self.assertGreaterEqual(check_sum, 220)