Spaces:
Sleeping
Sleeping
| import os | |
| from scipy.io import loadmat | |
| from menpo.shape.pointcloud import PointCloud | |
| from menpo.transform import ThinPlateSplines | |
| import menpo.transform as mt | |
| import menpo.io as mio | |
| from glob import glob | |
| from thirdparty.face_of_art.deformation_functions import * | |
| # landmark indices by facial feature | |
| jaw_indices = np.arange(0, 17) | |
| lbrow_indices = np.arange(17, 22) | |
| rbrow_indices = np.arange(22, 27) | |
| upper_nose_indices = np.arange(27, 31) | |
| lower_nose_indices = np.arange(31, 36) | |
| leye_indices = np.arange(36, 42) | |
| reye_indices = np.arange(42, 48) | |
| outer_mouth_indices = np.arange(48, 60) | |
| inner_mouth_indices = np.arange(60, 68) | |
| # flipped landmark indices | |
| mirrored_parts_68 = np.hstack([ | |
| jaw_indices[::-1], rbrow_indices[::-1], lbrow_indices[::-1], | |
| upper_nose_indices, lower_nose_indices[::-1], | |
| np.roll(reye_indices[::-1], 4), np.roll(leye_indices[::-1], 4), | |
| np.roll(outer_mouth_indices[::-1], 7), | |
| np.roll(inner_mouth_indices[::-1], 5) | |
| ]) | |
| def load_bb_files(bb_file_dirs): | |
| """load bounding box mat file for challenging, common, full & training datasets""" | |
| bb_files_dict = {} | |
| for bb_file in bb_file_dirs: | |
| bb_mat = loadmat(bb_file)['bounding_boxes'] | |
| num_imgs = np.max(bb_mat.shape) | |
| for i in range(num_imgs): | |
| name = bb_mat[0][i][0][0][0][0] | |
| bb_init = bb_mat[0][i][0][0][1] - 1 # matlab indicies | |
| bb_gt = bb_mat[0][i][0][0][2] - 1 # matlab indicies | |
| if str(name) in bb_files_dict.keys(): | |
| print (str(name) + ' already exists') | |
| else: | |
| bb_files_dict[str(name)] = (bb_init, bb_gt) | |
| return bb_files_dict | |
| def load_bb_dictionary(bb_dir, mode, test_data='full'): | |
| """create bounding box dictionary of input dataset: train/common/full/challenging""" | |
| if mode == 'TRAIN': | |
| bb_dirs = \ | |
| ['bounding_boxes_afw.mat', 'bounding_boxes_helen_trainset.mat', 'bounding_boxes_lfpw_trainset.mat'] | |
| else: | |
| if test_data == 'common': | |
| bb_dirs = \ | |
| ['bounding_boxes_helen_testset.mat', 'bounding_boxes_lfpw_testset.mat'] | |
| elif test_data == 'challenging': | |
| bb_dirs = ['bounding_boxes_ibug.mat'] | |
| elif test_data == 'full': | |
| bb_dirs = \ | |
| ['bounding_boxes_ibug.mat', 'bounding_boxes_helen_testset.mat', 'bounding_boxes_lfpw_testset.mat'] | |
| elif test_data == 'training': | |
| bb_dirs = \ | |
| ['bounding_boxes_afw.mat', 'bounding_boxes_helen_trainset.mat', 'bounding_boxes_lfpw_trainset.mat'] | |
| else: | |
| bb_dirs = None | |
| if mode == 'TEST' and test_data not in ['full', 'challenging', 'common', 'training']: | |
| bb_files_dict = None | |
| else: | |
| bb_dirs = [os.path.join(bb_dir, dataset) for dataset in bb_dirs] | |
| bb_files_dict = load_bb_files(bb_dirs) | |
| return bb_files_dict | |
| def center_margin_bb(bb, img_bounds, margin=0.25): | |
| """create new bounding box with input margin""" | |
| bb_size = ([bb[0, 2] - bb[0, 0], bb[0, 3] - bb[0, 1]]) | |
| margins = (np.max(bb_size) * (1 + margin) - bb_size) / 2 | |
| bb_new = np.zeros_like(bb) | |
| bb_new[0, 0] = np.maximum(bb[0, 0] - margins[0], 0) | |
| bb_new[0, 2] = np.minimum(bb[0, 2] + margins[0], img_bounds[1]) | |
| bb_new[0, 1] = np.maximum(bb[0, 1] - margins[1], 0) | |
| bb_new[0, 3] = np.minimum(bb[0, 3] + margins[1], img_bounds[0]) | |
| return bb_new | |
| def crop_to_face_image(img, bb_dictionary=None, gt=True, margin=0.25, image_size=256, normalize=True, | |
| return_transform=False): | |
| """crop face image using bounding box dictionary, or GT landmarks""" | |
| name = img.path.name | |
| img_bounds = img.bounds()[1] | |
| # if there is no bounding-box dict and GT landmarks are available, use it to determine the bounding box | |
| if bb_dictionary is None and img.has_landmarks: | |
| grp_name = img.landmarks.group_labels[0] | |
| bb_menpo = img.landmarks[grp_name].bounding_box().points | |
| bb = np.array([[bb_menpo[0, 1], bb_menpo[0, 0], bb_menpo[2, 1], bb_menpo[2, 0]]]) | |
| elif bb_dictionary is not None: | |
| if gt: | |
| bb = bb_dictionary[name][1] # ground truth | |
| else: | |
| bb = bb_dictionary[name][0] # init from face detector | |
| else: | |
| bb = None | |
| if bb is not None: | |
| # add margin to bounding box | |
| bb = center_margin_bb(bb, img_bounds, margin=margin) | |
| bb_pointcloud = PointCloud(np.array([[bb[0, 1], bb[0, 0]], | |
| [bb[0, 3], bb[0, 0]], | |
| [bb[0, 3], bb[0, 2]], | |
| [bb[0, 1], bb[0, 2]]])) | |
| if return_transform: | |
| face_crop, bb_transform = img.crop_to_pointcloud(bb_pointcloud, return_transform=True) | |
| else: | |
| face_crop = img.crop_to_pointcloud(bb_pointcloud) | |
| else: | |
| # if there is no bounding box/gt landmarks, use entire image | |
| face_crop = img.copy() | |
| bb_transform = None | |
| # if face crop is not a square - pad borders with mean pixel value | |
| h, w = face_crop.shape | |
| diff = h - w | |
| if diff < 0: | |
| face_crop.pixels = np.pad(face_crop.pixels, ((0, 0), (0, -1 * diff), (0, 0)), 'mean') | |
| elif diff > 0: | |
| face_crop.pixels = np.pad(face_crop.pixels, ((0, 0), (0, 0), (0, diff)), 'mean') | |
| if return_transform: | |
| face_crop, rescale_transform = face_crop.resize([image_size, image_size], return_transform=True) | |
| if bb_transform is None: | |
| transform_chain = rescale_transform | |
| else: | |
| transform_chain = mt.TransformChain(transforms=(rescale_transform, bb_transform)) | |
| else: | |
| face_crop = face_crop.resize([image_size, image_size]) | |
| if face_crop.n_channels == 4: | |
| face_crop.pixels = face_crop.pixels[:3, :, :] | |
| if normalize: | |
| face_crop.pixels = face_crop.rescale_pixels(0., 1.).pixels | |
| if return_transform: | |
| return face_crop, transform_chain | |
| else: | |
| return face_crop | |
| def augment_face_image(img, image_size=256, crop_size=248, angle_range=30, flip=True): | |
| """basic image augmentation: random crop, rotation and horizontal flip""" | |
| # taken from MDM: https://github.com/trigeorgis/mdm | |
| def mirror_landmarks_68(lms, im_size): | |
| return PointCloud(abs(np.array([0, im_size[1]]) - lms.as_vector( | |
| ).reshape(-1, 2))[mirrored_parts_68]) | |
| # taken from MDM: https://github.com/trigeorgis/mdm | |
| def mirror_image(im): | |
| im = im.copy() | |
| im.pixels = im.pixels[..., ::-1].copy() | |
| for group in im.landmarks: | |
| lms = im.landmarks[group] | |
| if lms.points.shape[0] == 68: | |
| im.landmarks[group] = mirror_landmarks_68(lms, im.shape) | |
| return im | |
| flip_rand = np.random.random() > 0.5 | |
| # rot_rand = np.random.random() > 0.5 | |
| # crop_rand = np.random.random() > 0.5 | |
| rot_rand = True # like ECT: https://github.com/HongwenZhang/ECT-FaceAlignment | |
| crop_rand = True # like ECT: https://github.com/HongwenZhang/ECT-FaceAlignment | |
| if crop_rand: | |
| lim = image_size - crop_size | |
| min_crop_inds = np.random.randint(0, lim, 2) | |
| max_crop_inds = min_crop_inds + crop_size | |
| img = img.crop(min_crop_inds, max_crop_inds) | |
| if flip and flip_rand: | |
| img = mirror_image(img) | |
| if rot_rand: | |
| rot_angle = 2 * angle_range * np.random.random_sample() - angle_range | |
| img = img.rotate_ccw_about_centre(rot_angle) | |
| img = img.resize([image_size, image_size]) | |
| return img | |
| def augment_menpo_img_ns(img, img_dir_ns, p_ns=0.): | |
| """texture style image augmentation using stylized copies in *img_dir_ns*""" | |
| img = img.copy() | |
| if p_ns > 0.5: | |
| ns_augs = glob(os.path.join(img_dir_ns, img.path.name.split('.')[0] + '_ns*')) | |
| num_augs = len(ns_augs) | |
| if num_augs > 0: | |
| ns_ind = np.random.randint(0, num_augs) | |
| ns_aug = mio.import_image(ns_augs[ns_ind]) | |
| ns_pixels = ns_aug.pixels | |
| img.pixels = ns_pixels | |
| return img | |
| def augment_menpo_img_geom(img, p_geom=0.): | |
| """geometric style image augmentation using random face deformations""" | |
| img = img.copy() | |
| if p_geom > 0.5: | |
| grp_name = img.landmarks.group_labels[0] | |
| lms_geom_warp = deform_face_geometric_style(img.landmarks[grp_name].points.copy(), p_scale=p_geom, p_shift=p_geom) | |
| img = warp_face_image_tps(img, PointCloud(lms_geom_warp), grp_name) | |
| return img | |
| def warp_face_image_tps(img, new_shape, lms_grp_name='PTS', warp_mode='constant'): | |
| """warp image to new landmarks using TPS interpolation""" | |
| tps = ThinPlateSplines(new_shape, img.landmarks[lms_grp_name]) | |
| try: | |
| img_warp = img.warp_to_shape(img.shape, tps, mode=warp_mode) | |
| img_warp.landmarks[lms_grp_name] = new_shape | |
| return img_warp | |
| except np.linalg.linalg.LinAlgError as err: | |
| print ('Error:'+str(err)+'\nUsing original landmarks for:\n'+str(img.path)) | |
| return img | |
| def load_menpo_image_list( | |
| img_dir, train_crop_dir, img_dir_ns, mode, bb_dictionary=None, image_size=256, margin=0.25, | |
| bb_type='gt', test_data='full', augment_basic=True, augment_texture=False, p_texture=0, | |
| augment_geom=False, p_geom=0, verbose=False, return_transform=False): | |
| """load images from image dir to create menpo-type image list""" | |
| def crop_to_face_image_gt(img): | |
| return crop_to_face_image(img, bb_dictionary, gt=True, margin=margin, image_size=image_size, | |
| return_transform=return_transform) | |
| def crop_to_face_image_init(img): | |
| return crop_to_face_image(img, bb_dictionary, gt=False, margin=margin, image_size=image_size, | |
| return_transform=return_transform) | |
| def crop_to_face_image_test(img): | |
| return crop_to_face_image(img, bb_dictionary=None, margin=margin, image_size=image_size, | |
| return_transform=return_transform) | |
| def augment_menpo_img_ns_rand(img): | |
| return augment_menpo_img_ns(img, img_dir_ns, p_ns=1. * (np.random.rand() < p_texture)[0]) | |
| def augment_menpo_img_geom_rand(img): | |
| return augment_menpo_img_geom(img, p_geom=1. * (np.random.rand() < p_geom)[0]) | |
| if mode is 'TRAIN': | |
| if train_crop_dir is None: | |
| img_set_dir = os.path.join(img_dir, 'training') | |
| out_image_list = mio.import_images(img_set_dir, verbose=verbose, normalize=False) | |
| if bb_type is 'gt': | |
| out_image_list = out_image_list.map(crop_to_face_image_gt) | |
| elif bb_type is 'init': | |
| out_image_list = out_image_list.map(crop_to_face_image_init) | |
| else: | |
| img_set_dir = os.path.join(img_dir, train_crop_dir) | |
| out_image_list = mio.import_images(img_set_dir, verbose=verbose) | |
| # perform image augmentation | |
| if augment_texture and p_texture > 0: | |
| out_image_list = out_image_list.map(augment_menpo_img_ns_rand) | |
| if augment_geom and p_geom > 0: | |
| out_image_list = out_image_list.map(augment_menpo_img_geom_rand) | |
| if augment_basic: | |
| out_image_list = out_image_list.map(augment_face_image) | |
| else: # if mode is 'TEST', load test data | |
| if test_data in ['full', 'challenging', 'common', 'training', 'test']: | |
| img_set_dir = os.path.join(img_dir, test_data) | |
| out_image_list = mio.import_images(img_set_dir, verbose=verbose, normalize=False) | |
| if bb_type is 'gt': | |
| out_image_list = out_image_list.map(crop_to_face_image_gt) | |
| elif bb_type is 'init': | |
| out_image_list = out_image_list.map(crop_to_face_image_init) | |
| else: | |
| img_set_dir = os.path.join(img_dir, test_data+'*') | |
| out_image_list = mio.import_images(img_set_dir, verbose=verbose, normalize=False) | |
| out_image_list = out_image_list.map(crop_to_face_image_test) | |
| return out_image_list | |