Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,337 Bytes
1d24639 |
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 |
import os
import cv2
import numpy as np
from PIL import Image
from batch_face import RetinaFace
def _get_square_face(coord, image, padding_scale = 1.5):
x1, y1, x2, y2 = coord
# expand the face region by {padding_scale} times
length = ((x2 - x1) + (y2 - y1)) // 2
x1 = x1 - length * (padding_scale - 1.0)
x2 = x2 + length * (padding_scale - 1.0)
y1 = y1 - length * (padding_scale - 1.0)
y2 = y2 + length * (padding_scale - 1.0)
# Move the center upside a little
y1 -= length * (padding_scale - 1.0) * 0.2
y2 -= length * (padding_scale - 1.0) * 0.2
# get square image
center = (x1 + x2) // 2, (y1 + y2) // 2
length = max(x2 - x1, y2 - y1) // 2
x1 = max(int(round(center[0] - length)), 0)
x2 = min(int(round(center[0] + length)), image.shape[1])
y1 = max(int(round(center[1] - length)), 0)
y2 = min(int(round(center[1] + length)), image.shape[0])
return image[y1:y2, x1:x2]
def _get_face_coord(face_detector, frame_cv2):
faces = face_detector(frame_cv2, cv=True)
if len(faces) == 0:
raise ValueError("Face is not detected")
else:
coord = faces[0][0]
return coord
def _smooth_coord(last_coord, current_coord, smooth_factor=0.1):
change = np.array(current_coord) - np.array(last_coord)
# smooth the change to 0.1 times
change = change * smooth_factor
return (np.array(last_coord) + np.array(change)).astype(int).tolist()
def get_face_img(face_detector, input_frame_path):
print("Detecting face in the image...")
frame_cv2 = cv2.imread(input_frame_path)
coord = _get_face_coord(face_detector, frame_cv2)
face = _get_square_face(coord, frame_cv2)
face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
return Image.fromarray(face), coord
def get_faces_video(face_detector, input_video_path):
output_frames = []
output_coords = []
last_coord = None
print("Detecting faces in the video...")
cap = cv2.VideoCapture(input_video_path)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
face_coord = _get_face_coord(face_detector, frame)
if last_coord is not None:
face_coord = _smooth_coord(last_coord, face_coord)
last_coord = face_coord
face = _get_square_face(face_coord, frame)
face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
face_pil = Image.fromarray(face)
output_frames.append(face_pil)
output_coords.append(face_coord)
cap.release()
return output_frames, output_coords
if __name__ == '__main__':
import torch
face_detector = RetinaFace(gpu_id=0) if torch.cuda.is_available() else RetinaFace(gpu_id=-1)
# test for image
input_frame_path = './test_imgs/makeup/1.jpg'
face, _ = get_face_img(face_detector, input_frame_path)
face.save('face.png')
print("Image saved to face.png")
# test for video
import imageio
from tqdm import tqdm
frames, _ = get_faces_video(face_detector, './test_imgs/input_video.mp4')
print("Number of frames: ", len(frames))
writer = imageio.get_writer('face.mp4', fps=30, macro_block_size=1, quality=8, codec="libx264")
for frame in tqdm(frames):
writer.append_data(np.array(frame.resize((512, 512))))
writer.close()
print("Video saved to face.mp4")
|