File size: 3,827 Bytes
2e4e201 9a70da5 2e4e201 9a70da5 2e4e201 9a70da5 2e4e201 9a70da5 2e4e201 9a70da5 2e4e201 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import os
import sys
import numpy as np
import cv2
import time
from tqdm import tqdm
import multiprocessing
import glob
import mediapipe as mp
from mediapipe import solutions
from mediapipe.framework.formats import landmark_pb2
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from src.utils import face_landmark
CUR_DIR = os.path.dirname(__file__)
class LMKExtractor():
def __init__(self, FPS=25):
# Create an FaceLandmarker object.
self.mode = mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE
base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/face_landmarker_v2_with_blendshapes.task'))
base_options.delegate = mp.tasks.BaseOptions.Delegate.CPU
options = vision.FaceLandmarkerOptions(base_options=base_options,
running_mode=self.mode,
output_face_blendshapes=True,
output_facial_transformation_matrixes=True,
num_faces=1)
self.detector = face_landmark.FaceLandmarker.create_from_options(options)
self.last_ts = 0
self.frame_ms = int(1000 / FPS)
# det_base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/blaze_face_short_range.tflite'))
# det_options = vision.FaceDetectorOptions(base_options=det_base_options)
# self.det_detector = vision.FaceDetector.create_from_options(det_options)
def __call__(self, img):
frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
# t0 = time.time()
# if self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.VIDEO:
# det_result = self.det_detector.detect(image)
# if len(det_result.detections) != 1:
# return None
# self.last_ts += self.frame_ms
# try:
# detection_result, mesh3d = self.detector.detect_for_video(image, timestamp_ms=self.last_ts)
# except:
# return None
# elif self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE:
# det_result = self.det_detector.detect(image)
# if len(det_result.detections) != 1:
# return None
try:
detection_result, mesh3d = self.detector.detect(image)
except:
return None
bs_list = detection_result.face_blendshapes
if len(bs_list) == 1:
bs = bs_list[0]
bs_values = []
for index in range(len(bs)):
bs_values.append(bs[index].score)
bs_values = bs_values[1:] # remove neutral
trans_mat = detection_result.facial_transformation_matrixes[0]
face_landmarks_list = detection_result.face_landmarks
face_landmarks = face_landmarks_list[0]
lmks = []
for index in range(len(face_landmarks)):
x = face_landmarks[index].x
y = face_landmarks[index].y
z = face_landmarks[index].z
lmks.append([x, y, z])
lmks = np.array(lmks)
lmks3d = np.array(mesh3d.vertex_buffer)
lmks3d = lmks3d.reshape(-1, 5)[:, :3]
mp_tris = np.array(mesh3d.index_buffer).reshape(-1, 3) + 1
return {
"lmks": lmks,
'lmks3d': lmks3d,
"trans_mat": trans_mat,
'faces': mp_tris,
"bs": bs_values
}
else:
# print('multiple faces in the image: {}'.format(img_path))
return None
|