Spaces:
Paused
Paused
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 |