| | |
| | import argparse |
| | import os |
| | import random |
| | import numpy as np |
| | import cv2 |
| | import glob |
| | from mediapipe import solutions |
| | from mediapipe.framework.formats import landmark_pb2 |
| | import numpy as np |
| | import matplotlib.pyplot as plt |
| | import torch |
| | import torch.nn as nn |
| | from scipy.signal import savgol_filter |
| | import onnxruntime as ort |
| | from collections import OrderedDict |
| | import mediapipe as mp |
| | from mediapipe.tasks import python |
| | from mediapipe.tasks.python import vision |
| |
|
| |
|
| | from tqdm import tqdm |
| |
|
| |
|
| | def infer_bs(root_path): |
| | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| | base_options = python.BaseOptions(model_asset_path="./data_utils/blendshape_capture/face_landmarker.task") |
| | options = vision.FaceLandmarkerOptions(base_options=base_options, |
| | output_face_blendshapes=True, |
| | output_facial_transformation_matrixes=True, |
| | num_faces=1) |
| | detector = vision.FaceLandmarker.create_from_options(options) |
| |
|
| | for i in os.listdir(root_path): |
| | if i.endswith(".mp4"): |
| | mp4_path = os.path.join(root_path, i) |
| | npy_path = os.path.join(root_path, "bs.npy") |
| | if os.path.exists(npy_path): |
| | print("npy file exists:", i.split(".")[0]) |
| | continue |
| | else: |
| | print("npy file not exists:", i.split(".")[0]) |
| | image_path = os.path.join(root_path, "img/temp.png") |
| | os.makedirs(os.path.join(root_path, 'img/'), exist_ok=True) |
| | cap = cv2.VideoCapture(mp4_path) |
| | fps = cap.get(cv2.CAP_PROP_FPS) |
| | frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| | print("fps:", fps) |
| | print("frame_count:", frame_count) |
| | k = 0 |
| | total = frame_count |
| | bs = np.zeros((int(total), 52), dtype=np.float32) |
| | print("total:", total) |
| | print("videoPath:{} fps:{},k".format(mp4_path.split('/')[-1], fps)) |
| | pbar = tqdm(total=int(total)) |
| | while (cap.isOpened()): |
| | ret, frame = cap.read() |
| | if ret: |
| | cv2.imwrite(image_path, frame) |
| | image = mp.Image.create_from_file(image_path) |
| | result = detector.detect(image) |
| | face_blendshapes_scores = [face_blendshapes_category.score for face_blendshapes_category in |
| | result.face_blendshapes[0]] |
| | blendshape_coef = np.array(face_blendshapes_scores)[1:] |
| | blendshape_coef = np.append(blendshape_coef, 0) |
| | bs[k] = blendshape_coef |
| | pbar.update(1) |
| | k += 1 |
| | else: |
| | break |
| | cap.release() |
| | pbar.close() |
| | |
| | |
| | output = np.zeros((bs.shape[0], bs.shape[1])) |
| | for j in range(bs.shape[1]): |
| | output[:, j] = savgol_filter(bs[:, j], 5, 3) |
| | np.save(npy_path, output) |
| | print(np.shape(output)) |
| |
|
| |
|
| | if __name__ == '__main__': |
| | parser = argparse.ArgumentParser() |
| | parser.add_argument("--path", type=str, help="idname of target person") |
| | args = parser.parse_args() |
| | infer_bs(args.path) |
| |
|
