xyz / app.py
andro-flock's picture
Upload folder using huggingface_hub
0d66750 verified
raw
history blame
1.17 kB
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