Upload folder using huggingface_hub

app.py

@@ -12,21 +12,21 @@ def prepare_image(img_path):
     body_bboxes, body_keypoints = pose_result.boxes, pose_result.keypoints 
     # generate heatmap
     heatmap = get_heatmap(img_path)
-    heatmap_pil = Image.fromarray(heatmap)
+    heatmap_pil = Image.fromarray(heatmap).convert('RGB')
     # pose on heatmap
-    heatmap_n_pose = draw_pose(heatmap_pil,body_keypoints) #PIL.Image
+    heatmap_n_pose = draw_pose(heatmap_pil,body_keypoints,0.5) #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
+    face_y1, face_h = face_bboxes[0][1], face_bboxes[0][2] #xywh of first face
+    face_y2 = face_y1 + face_h
+    body_box = body_bboxes.numpy().xyxy[0]
+    body_box[1] = face_y2 #shift below head
     
     # 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))
+    x1, y1, x2, y2 = map(int, body_box)
+    cropped_region = heatmap_n_pose.crop((x1, y1, x2, y2))
     original_image = Image.open(img_path)
-    original_image.paste(cropped_region, (x, y))
+    original_image.paste(cropped_region, (x1, y1))
     
     # Return the modified image
     return original_image