Add command line arguments, and general code cleanup.

This commit is contained in:
Nathan Vegdahl 2022-05-31 10:01:54 -07:00
parent e4312b6ab8
commit 353b483e08

View File

@ -1,32 +1,54 @@
#!/usr/bin/env python3
import argparse
from struct import pack
import numpy as np
import cv2 as cv
import mediapipe as mp
if __name__ == "__main__":
mp_face_mesh = mp.solutions.face_mesh
arg_parser = argparse.ArgumentParser(description=
"""
Processes video of a human face into 3D face mesh animation
data. Currently only supports video files with a single face.
"""
)
arg_parser.add_argument("input_video", help="The input video file with a person's face in it.")
arg_parser.add_argument("output_mdd", help="The .mdd file to write the mesh animation data to.")
args = arg_parser.parse_args()
# Compute the mesh points. One mesh per frame.
meshes = []
video_path = args.input_video
mdd_path = args.output_mdd
# Compute the mesh points from the input video.
meshes = [] # One mesh per frame.
point_count = 0
width = 0
height = 0
fps = 0
aspect_ratio = 1.0
with mp_face_mesh.FaceMesh(
with mp.solutions.face_mesh.FaceMesh(
static_image_mode=False, # Set false for video.
max_num_faces=1,
refine_landmarks=True,
min_detection_confidence=0.5) as face_mesh:
video = cv.VideoCapture(video_path)
width = video.get(cv.CAP_PROP_FRAME_WIDTH)
height = video.get(cv.CAP_PROP_FRAME_HEIGHT)
fps = video.get(cv.CAP_PROP_FPS)
aspect_ratio = width / height
print("Input video fps: {}".format(fps))
print("Input video resolution: {}x{}".format(int(width), int(height)))
i = 1
cap = cv.VideoCapture("videos/BDGFACE 5_12.mp4")
while cap.isOpened():
while video.isOpened():
ret, image = video.read()
print("\rReading frame", i, end = "")
ret, image = cap.read()
if not ret:
break
aspect_ratio = image.shape[1] / image.shape[0]
results = face_mesh.process(cv.cvtColor(image, cv.COLOR_BGR2RGB))
if not results.multi_face_landmarks:
meshes += [None]
@ -35,15 +57,14 @@ if __name__ == "__main__":
meshes += [results.multi_face_landmarks[0].landmark]
i += 1
cap.release()
video.release()
print("\rRead {} frames. ".format(len(meshes)))
print("Generated vert count:", point_count)
# Write the mdd file.
fps = 24.0 # TODO: get this from the video file.
frame_count = len(meshes)
if frame_count > 0:
with open("test.mdd", 'wb') as mdd:
print("Vert count:", point_count)
with open(mdd_path, 'wb') as mdd:
mdd.write(pack(">2i", frame_count, point_count))
mdd.write(pack(">%df" % (frame_count), *[frame / fps for frame in range(frame_count)]))
@ -59,5 +80,5 @@ if __name__ == "__main__":
for point in mesh:
mdd.write(pack(">3f", point.x, point.y / aspect_ratio, point.z))
i += 1
print()
print("\rWrote {} frames. ".format(len(meshes)))