Last active
October 6, 2023 15:00
-
-
Save recoilme/1ada5bcdd62636b2198b930fd90005a9 to your computer and use it in GitHub Desktop.
videomaskbyface.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import cv2 | |
import os | |
import time | |
import mediapipe as mp | |
import numpy as np | |
def extract_faces(image): | |
center_faces = [] | |
BaseOptions = mp.tasks.BaseOptions | |
FaceDetector = mp.tasks.vision.FaceDetector | |
FaceDetectorOptions = mp.tasks.vision.FaceDetectorOptions | |
VisionRunningMode = mp.tasks.vision.RunningMode | |
# Create a face detector instance with the image mode: | |
options = FaceDetectorOptions( | |
base_options=BaseOptions(model_asset_path='blaze_face_short_range.tflite'), | |
running_mode=VisionRunningMode.IMAGE) | |
detector = FaceDetector.create_from_options(options) | |
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image) | |
# The face detector must be created with the image mode. | |
face_detector_result = detector.detect(mp_image) | |
#print(face_detector_result) | |
for detection in face_detector_result.detections: | |
bbox = detection.bounding_box | |
center_faces.append(([int(bbox.origin_x+bbox.width/2), int(bbox.origin_y+bbox.height/2)])) | |
return center_faces | |
def find_normalized_distances(image, face_coordinates, low, max): | |
height, width, _ = image.shape | |
distances = np.full((height, width), np.inf) | |
for (center_x, center_y) in face_coordinates: | |
distances_old = distances.copy() | |
distances = np.sqrt((np.arange(width) - center_x)**2 + (np.arange(height)[:,np.newaxis] - center_y)**2) | |
distances[distances < 0.0] = 0.0 | |
distances = np.minimum(distances, distances_old) | |
if np.max(distances) == np.inf: | |
return np.zeros((height, width)) | |
# Normalize the distance values between min and max | |
OldRange = (np.max(distances) - np.min(distances)) | |
NewRange = (max - low) | |
normalized_distances = (((distances - np.min(distances)) * NewRange) / OldRange) + low | |
return normalized_distances | |
def merge_masks_with_image(base_image, normalized_distances): | |
b_channel, g_channel, r_channel = cv2.split(base_image) | |
alpha_gradient = (255 - normalized_distances * 255).astype(np.uint8) | |
img_BGRA = cv2.merge((b_channel, g_channel, r_channel, alpha_gradient)) | |
return img_BGRA | |
# Open the video file | |
video = cv2.VideoCapture('input.mp4') | |
# Get the video frame rate and size | |
fps = int(video.get(cv2.CAP_PROP_FPS)) | |
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
# Create a directory to store the extracted frames | |
if not os.path.exists('frames'): | |
os.makedirs('frames') | |
# Initialize a counter for the frames | |
frame_count = 0 | |
# Loop through the video frames | |
start_time = time.time() | |
while True: | |
# Read the next frame from the video | |
ret, frame = video.read() | |
# If there are no more frames, break out of the loop | |
if not ret: | |
break | |
# Extract faces | |
face_coordinates = extract_faces(frame) | |
# Create transparency mask | |
normalized_distances = find_normalized_distances(frame, face_coordinates, 0, 1.0) | |
alpha_gradient = ((normalized_distances) * 255).astype(np.uint8) | |
# Normalize the depth map to the range [0, 255] | |
depth_map_normalized = cv2.normalize(alpha_gradient, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U) | |
cv2.imwrite(f'frames/frame_{frame_count}.png',depth_map_normalized) | |
#image_with_alpha = merge_masks_with_image(frame, normalized_distances) | |
#cv2.imwrite("result_image.png", image_with_alpha) | |
#cv2.imwrite(f'frames/frame_{frame_count}.png', cv2.merge((mask[:, :, 0:3], alpha_gradient))) | |
# Increment the frame counter | |
frame_count += 1 | |
# Release the video file | |
video.release() | |
end_time = time.time() | |
print("Running time:", end_time - start_time, "seconds") | |
# Create a new video file from the extracted frames | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
#image = cv2.imread("1.jpg") | |
#height, width, _ = image.shape | |
out = cv2.VideoWriter('output.mp4', fourcc, fps, (width, height)) | |
# Loop through the extracted frames and add them to the new video file | |
#frame_count = 2 | |
for i in range(frame_count): | |
#frame = cv2.imread("1.jpg") | |
frame = cv2.imread(f'frames/frame_{i}.png') | |
out.write(frame) | |
# Release the new video file | |
out.release() | |
# Remove the extracted frames directory and its contents | |
for file in os.listdir('frames'): | |
os.remove(os.path.join('frames', file)) | |
os.rmdir('frames') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment