import os
import cv2
from tqdm import tqdm
import face_recognition

# Input and output paths
frame_dir = 'your path'
process_dir = 'your path'

# Create output directory
os.makedirs(process_dir, exist_ok=True)

def expand_bbox(top, right, bottom, left, img_height, img_width, margin_ratio=0.25):
    """
    Expand face bounding box coordinates by adding margins.
    """
    height = bottom - top
    width = right - left
    margin_h = int(height * margin_ratio)
    margin_w = int(width * margin_ratio)

    # Expand boundaries
    new_top = max(0, top - margin_h)
    new_bottom = min(img_height, bottom + margin_h)
    new_left = max(0, left - margin_w)
    new_right = min(img_width, right + margin_w)

    return new_top, new_right, new_bottom, new_left

# Iterate through all subdirectories in frame_dir
for video_name in os.listdir(frame_dir):
    video_path = os.path.join(frame_dir, video_name)
    if not os.path.isdir(video_path):
        continue

    # Create corresponding video output directory
    output_video_dir = os.path.join(process_dir, video_name)
    os.makedirs(output_video_dir, exist_ok=True)

    # Store face coordinates of the first frame
    first_face_coords = None

    # Process each frame
    for frame_name in tqdm(os.listdir(video_path), desc=f"Processing {video_name}"):
        frame_path = os.path.join(video_path, frame_name)
        output_frame_path = os.path.join(output_video_dir, frame_name)

        # Read frame
        frame = cv2.imread(frame_path)
        if frame is None:
            print(f"Failed to read frame: {frame_path}")
            continue

        img_height, img_width = frame.shape[:2]

        # If first frame's face coordinates not recorded, attempt detection
        if first_face_coords is None:
            face_locations = face_recognition.face_locations(frame)
            if face_locations:
                # Use coordinates of the first detected face
                first_face_coords = face_locations[0]
            else:
                print(f"No face detected in: {frame_path}")
                continue

        # If we have the first frame's face coordinates, use them
        if first_face_coords:
            top, right, bottom, left = first_face_coords

            # Expand bounding box
            new_top, new_right, new_bottom, new_left = expand_bbox(
                top, right, bottom, left, img_height, img_width, margin_ratio=0.10
            )

            # Crop the expanded region
            face_image = frame[new_top:new_bottom, new_left:new_right]

            # Resize face to 256x256
            resized_face = cv2.resize(face_image, (256, 256))

            # Save to target directory
            cv2.imwrite(output_frame_path, resized_face)

print("Processing completed!")