from PIL import Image
from face_detection import detect_faces
from heatmap import get_heatmap
from pose import detect_pose, draw_pose

def prepare_image(img_path):
    # detect face
    annotated_image, face_bboxes = detect_faces(img_path)
    # detect pose + bounding_box
    pose_results = detect_pose(img_path)
    pose_result = pose_results[0]
    body_bboxes, body_keypoints = pose_result.boxes, pose_result.keypoints 
    # generate heatmap
    heatmap = get_heatmap(img_path)
    heatmap_pil = Image.fromarray(heatmap)
    # pose on heatmap
    heatmap_n_pose = draw_pose(heatmap_pil,body_keypoints) #PIL.Image
    
    # area in bounding_box below head (rectangular)
    face_y1, face_w = face_bboxes[0][1] #xywh
    face_center_y = (face_w - face_y1) // 2
    body_box = body_bboxes.numpy().xywh[0]
    body_box[1] = face_center_y
    
    # Calculate the coordinates for cropping
    x, y, w, h = map(int, body_box)
    cropped_region = heatmap_n_pose.crop((x, y, x + w, y + h))
    original_image = Image.open(img_path)
    original_image.paste(cropped_region, (x, y))
    
    # Return the modified image
    return original_image