videomaskbyface.py

videomaskbyface.py

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')