Spaces:
Sleeping
Sleeping
File size: 2,124 Bytes
40ac571 |
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 |
import os
import numpy as np
import imageio
import cv2
from tqdm import tqdm
from media_pipe import FaceMeshDetector, FaceMeshAlign
def process_video_with_align(video_path, save_dir):
face_detector = FaceMeshDetector()
face_aligner = FaceMeshAlign()
frames = imageio.get_reader(video_path)
face_results = []
motions = []
# Process first frame to get reference
first_frame = next(iter(frames))
first_frame_rgb = np.array(first_frame)
motion, ref_result = face_detector(first_frame_rgb)
if ref_result is None:
print("No face detected in the first frame. Exiting.")
return
face_results.append(ref_result)
motions.append(motion)
# Process remaining frames
for frame in tqdm(frames):
frame_rgb = np.array(frame)
motion, face_result = face_detector(frame_rgb)
if face_result is None:
continue
face_results.append(face_result)
motions.append(motion)
# Perform alignment
aligned_motions = face_aligner(ref_result, face_results)
base_name = os.path.splitext(os.path.basename(video_path))[0]
# Save original results
npy_path = os.path.join(save_dir, f"{base_name}_mppose.npy")
np.save(npy_path, face_results)
gif_path = os.path.join(save_dir, f"{base_name}_mppose.gif")
imageio.mimsave(gif_path, motions, 'GIF', duration=0.2, loop=0)
# Save aligned results
aligned_npy_path = os.path.join(save_dir, f"{base_name}_mppose_aligned.npy")
np.save(aligned_npy_path, aligned_motions)
aligned_gif_path = os.path.join(save_dir, f"{base_name}_mppose_aligned.gif")
imageio.mimsave(aligned_gif_path, aligned_motions, 'GIF', duration=0.2, loop=0)
return npy_path, aligned_npy_path, len(face_results)
# Пример использования
video_path = "./test/test.mp4"
save_dir = "./test"
original_npy, aligned_npy, frame_count = process_video_with_align(video_path, save_dir)
print(f"Processed {frame_count} frames.")
print(f"Original NPY saved at: {original_npy}")
print(f"Aligned NPY saved at: {aligned_npy}")
print(f"GIFs saved in: {save_dir}")
|