Spaces:
Sleeping
Sleeping
| import torch | |
| import os | |
| import sys | |
| import tqdm | |
| import glob | |
| import numpy as np | |
| import cv2 | |
| import face_alignment | |
| from skimage import io | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--data_source', type=str, default='./data/input') | |
| args = parser.parse_args() | |
| DATA_SOURCE = args.data_source | |
| device = torch.device('cuda:0') | |
| fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.THREE_D, flip_input=False, face_detector='blazeface') | |
| # DATA_SOURCE = '../../data/face_data' | |
| # DATA_SOURCE = '../../data/face_data' | |
| data_folder = os.path.join(DATA_SOURCE, 'images') | |
| frame_folders = sorted(glob.glob(data_folder + '/*')) | |
| for frame_folder in tqdm.tqdm(frame_folders): | |
| if 'background' in frame_folder: | |
| continue | |
| image_paths = glob.glob(frame_folder + '/image_*') | |
| images = np.stack([io.imread(image_path) for image_path in image_paths]) | |
| images = torch.from_numpy(images).float().permute(0, 3, 1, 2).to(device) | |
| results = fa.get_landmarks_from_batch(images, return_landmark_score=True) | |
| for i in range(len(results[0])): | |
| if results[1][i] is None: | |
| results[0][i] = np.zeros([68, 3], dtype=np.float32) | |
| results[1][i] = [np.zeros([68], dtype=np.float32)] | |
| if len(results[1][i]) > 1: | |
| total_score = 0.0 | |
| for j in range(len(results[1][i])): | |
| if np.sum(results[1][i][j]) > total_score: | |
| total_score = np.sum(results[1][i][j]) | |
| landmarks_i = results[0][i][j*68:(j+1)*68] | |
| scores_i = results[1][i][j:j+1] | |
| results[0][i] = landmarks_i | |
| results[1][i] = scores_i | |
| landmarks = np.concatenate([np.stack(results[0])[:, :, :2], np.stack(results[1]).transpose(0, 2, 1)], -1) | |
| i = 0 | |
| for image_path in image_paths: | |
| landmarks_path = image_path.replace('image_', 'landmarks_').replace('.jpg', '.npy') | |
| np.save(landmarks_path, landmarks[i]) | |
| i += 1 | |