Spaces:
Runtime error
Runtime error
File size: 1,972 Bytes
ec9a6bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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
|