import cv2
import torch
from PIL import Image
import os
from mtcnn import MTCNN

# Load the pre-trained YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')

# Initialize the MTCNN face detection model
mtcnn = MTCNN()

def detect_faces(image):
    # Convert image to RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # Detect faces in the image
    faces = mtcnn.detect_faces(image_rgb)
    return faces

def crop_face(image, face_coords, expansion_percentage):
    x, y, w, h = face_coords

    # Calculate expansion values
    expansion_x = int((w * expansion_percentage) / 100)
    expansion_y = int((h * expansion_percentage) / 100)

    # Adjust bounding box coordinates with expansion
    x -= expansion_x
    y -= expansion_y
    w += 2 * expansion_x
    h += 2 * expansion_y

    # Make sure the coordinates are within image boundaries
    x = max(0, x)
    y = max(0, y)
    w = min(image.shape[1] - 1, x + w) - x
    h = min(image.shape[0] - 1, y + h) - y

    # Crop the image around the expanded face region
    cropped_img = image[y:y + h, x:x + w]
    return cropped_img

def auto_zoom(input_dir, output_base_dir, expansion_percentage):
    # Loop over all files in the input directory
    for filename in os.listdir(input_dir):
        # Create the full input path and read the file
        input_path = os.path.join(input_dir, filename)
        img = cv2.imread(input_path)

        if img is None:
            continue

        # Run the image through the object detection model
        results = model(img)

        # Find the first human detected in the image
        human = next((x for x in results.xyxy[0] if int(x[5]) == 0), None)

        if human is None:
            print(f"No human detected in the image {input_path}.")
            os.remove(input_path)
            continue

        # Get the bounding box coordinates of the detected face
        face_coords = detect_faces(img)

        if len(face_coords) == 0:
            print(f"No face detected in the image {input_path}.")
            os.remove(input_path)
            continue

        # Crop the image around the expanded face region
        cropped_img = crop_face(img, face_coords[0]['box'], expansion_percentage)

        # Convert BGR image to RGB for PIL
        cropped_img_rgb = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)

        # Convert array to Image for visualization
        result_img = Image.fromarray(cropped_img_rgb)

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

        # Create the full output path and save the file
        output_path = os.path.join(output_dir, filename)
        result_img.save(output_path)

# Usage with 20% expansion of face while cropping - you can try different percentage
auto_zoom(r"F:\auto_crop\org_imgs", r"F:\auto_crop\new_images",20)