Spaces:
Sleeping
Sleeping
File size: 5,118 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 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import os
import numpy as np
import imageio
import cv2
from tqdm import tqdm
from media_pipe import FaceMeshDetector, FaceMeshAlign
from datetime import datetime
def crop_image_to_face(image, face_result, target_size):
height, width = image.shape[:2]
target_height, target_width = target_size
# Получаем координаты лица
landmarks = face_result['lmks']
min_x, min_y = np.min(landmarks, axis=0)[:2]
max_x, max_y = np.max(landmarks, axis=0)[:2]
# Вычисляем центр лица
center_x = (min_x + max_x) / 2 * width
center_y = (min_y + max_y) / 2 * height
# Вычисляем координаты для обрезки
left = int(max(center_x - target_width / 2, 0))
top = int(max(center_y - target_height / 2, 0))
right = int(min(left + target_width, width))
bottom = int(min(top + target_height, height))
# Обрезаем изображение
cropped_image = image[top:bottom, left:right]
# Если обрезанное изображение меньше целевого размера, добавляем отступы
if cropped_image.shape[0] < target_height or cropped_image.shape[1] < target_width:
pad_top = max(0, (target_height - cropped_image.shape[0]) // 2)
pad_bottom = max(0, target_height - cropped_image.shape[0] - pad_top)
pad_left = max(0, (target_width - cropped_image.shape[1]) // 2)
pad_right = max(0, target_width - cropped_image.shape[1] - pad_left)
cropped_image = cv2.copyMakeBorder(cropped_image, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT)
return cv2.resize(cropped_image, (target_width, target_height))
def process_image(image_path, face_detector, target_size):
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
motion, face_result = face_detector(image_rgb)
if face_result is not None:
cropped_image = crop_image_to_face(image_rgb, face_result, target_size)
cropped_motion, cropped_face_result = face_detector(cropped_image)
else:
cropped_image = None
cropped_motion = None
cropped_face_result = None
return motion, face_result, cropped_image, cropped_motion, cropped_face_result
def process_video_and_image(video_path, image_path, output_dir):
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
save_dir = os.path.join(output_dir, f"test_{timestamp}")
os.makedirs(save_dir, exist_ok=True)
face_detector = FaceMeshDetector()
face_aligner = FaceMeshAlign()
# Получаем размер кадра видео
video = cv2.VideoCapture(video_path)
frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
video.release()
# Process image
_, _, cropped_image, cropped_motion, cropped_face_result = process_image(image_path, face_detector, (frame_height, frame_width))
if cropped_face_result is None:
print("No face detected in the image. Exiting.")
return
# Save processed and cropped image
image_base_name = os.path.splitext(os.path.basename(image_path))[0]
processed_image_path = os.path.join(save_dir, f"{image_base_name}_processed.png")
cv2.imwrite(processed_image_path, cv2.cvtColor(cropped_motion, cv2.COLOR_RGB2BGR))
cropped_image_path = os.path.join(save_dir, f"{image_base_name}_cropped.png")
cv2.imwrite(cropped_image_path, cv2.cvtColor(cropped_image, cv2.COLOR_RGB2BGR))
# Process video
frames = imageio.get_reader(video_path)
face_results = []
motions = []
for frame in tqdm(frames, desc="Processing video"):
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(cropped_face_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)
print(f"Processed image saved at: {processed_image_path}")
print(f"Cropped image saved at: {cropped_image_path}")
print(f"Original NPY saved at: {npy_path}")
print(f"Aligned NPY saved at: {aligned_npy_path}")
print(f"GIFs saved in: {save_dir}")
# Пример использования
video_path = "./test/test.mp4"
image_path = "./test/123.png"
output_dir = "./test"
process_video_and_image(video_path, image_path, output_dir) |