Skip to content

Frames from video to Cropped Frames

qingy1337 edited this page Aug 30, 2024 · 1 revision
import os
import cv2
from mtcnn import MTCNN
from PIL import Image
from IPython.display import clear_output

# Initialize MTCNN for face detection
detector = MTCNN()

# Define the input and output directories
input_dir = "train_data/Zhi2"
output_dir = "train_data/Zhi"

# Create the output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)

# Process each image in the input directory
for filename in os.listdir(input_dir):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        image_path = os.path.join(input_dir, filename)
        img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)

        # Detect faces in the image
        faces = detector.detect_faces(img)

        # If a face is detected, crop and resize
        if faces:
            x, y, width, height = faces[0]['box']

            # Determine the size of the square based on the larger dimension of the face
            size = max(width, height)

            # Calculate the square crop coordinates
            crop_x1 = max(x + width // 2 - size // 2, 0)
            crop_y1 = max(y + height // 2 - size // 2, 0)
            crop_x2 = crop_x1 + size
            crop_y2 = crop_y1 + size

            # Crop the image to the square
            face_crop = img[crop_y1:crop_y2, crop_x1:crop_x2]

            # Resize the cropped image to 224x224
            face_resized = cv2.resize(face_crop, (224, 224))

            # Save the cropped and resized image
            save_path = os.path.join(output_dir, filename)
            Image.fromarray(face_resized).save(save_path)
        clear_output(wait=True)

print("Processing complete!")

Clone this wiki locally