output pose coords

app.py

@@ -26,6 +26,7 @@ with gr.Blocks() as demo:
             animation = gr.Video(label="Result")
             frames = gr.Gallery(type="pil", label="Frames", format="png")
             frames_thumb = gr.Gallery(type="pil", label="Thumbnails", format="png")
+            pose_coords = gr.JSON(label="Pose Coordinates")
 
     submit_btn.click(
         run_app, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, remove_bg, resize_inputs], outputs=[animation, frames]
@@ -36,7 +37,7 @@ with gr.Blocks() as demo:
     )
     
     inference_btn.click(
-        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, img_width, img_height, remove_bg, resize_inputs], outputs=[animation, frames, frames_thumb]
+        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, img_width, img_height, remove_bg, resize_inputs], outputs=[animation, frames, frames_thumb, pose_coords]
     )
 
 

main.py

@@ -58,7 +58,7 @@ import uuid
 import gc
 from numba import cuda
 import requests
-import uuid
+import json
 
 from huggingface_hub import hf_hub_download, HfApi
 
@@ -221,7 +221,7 @@ def get_pose(img, dwpose, outfile, crop=False):
         out_img = out_img.crop(bbox)
         out_img = ImageOps.expand(out_img, border=int(out_img.width*0.2), fill=(0,0,0))
     
-    return out_img
+    return out_img, pose
 
 
 def extract_frames(video_path, fps):
@@ -272,12 +272,13 @@ def prepare_inputs_train(images, bg_remove, dwpose, rembg_session):
         images = [removebg(img, rembg_session) for img in images]
 
     in_img = images[0]
-    in_pose = get_pose(in_img, dwpose, "in_pose.png")
+    in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")
     train_poses = []
     train_imgs = [resize_and_pad(img, in_img) for img in images[1:]]
     
     for i, img in enumerate(train_imgs):
-        train_poses.append(get_pose(img, dwpose, "tr_pose"+str(i)+".png"))
+        train_pose, _ = get_pose(img, dwpose, "tr_pose"+str(i)+".png")
+        train_poses.append(train_pose)
         
     return in_img, in_pose, train_imgs, train_poses
         
@@ -287,7 +288,7 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remo
 
     print("prepare_inputs_inference")
     
-    in_pose = get_pose(in_img, dwpose, "in_pose.png")
+    in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")
     
     frames = extract_frames(in_vid, fps)
     print("remove background", bg_remove)
@@ -302,14 +303,21 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remo
 
     progress_bar = tqdm(range(len(frames)), initial=0, desc="Frames")
     target_poses = []
+    target_poses_coords = []
     max_left = max_top = 999999
     max_right = max_bottom = 0
     it = frames
     if is_app:
         it = progress.tqdm(frames, desc="Pose Detection")
     for f in it:
-        tpose = get_pose(f, dwpose, "tar_pose"+str(len(target_poses))+".png")
+        tpose, tpose_coords = get_pose(f, dwpose, "tar_pose"+str(len(target_poses))+".png")
+        #print(tpose_coords)
+        coords = {}
+        for k in tpose_coords:
+            coords[k] = tpose_coords[k].tolist()
+        #print(coords)
         target_poses.append(tpose)
+        target_poses_coords.append(json.dumps(coords))
         progress_bar.update(1)
         
         bbox = tpose.getbbox()
@@ -332,14 +340,14 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remo
             tpose.save("out/"+"tar_pose"+str(len(target_poses_cropped))+".png")
         target_poses_cropped.append(tpose)
         
-    return in_img, target_poses_cropped, in_pose
+    return in_img, target_poses_cropped, in_pose, target_poses_coords
 
 
 def prepare_inputs(images, in_vid, fps, bg_remove, dwpose, rembg_session, resize='target', is_app=False):
     
     in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
     
-    in_img, target_poses_cropped, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize, is_app)
+    in_img, target_poses_cropped, _, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize, is_app)
     
     
     return in_img, in_pose, train_imgs, train_poses, target_poses_cropped
@@ -1125,7 +1133,7 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     images = [img[0] for img in images]
     in_img = images[0]
 
-    in_img, target_poses, in_pose = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', is_app)
+    in_img, target_poses, in_pose, target_poses_coords = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', is_app)
     
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
     #urls = save_temp_imgs(results)
@@ -1143,7 +1151,7 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     
     print("Done!")
     
-    return out_vid+'.webm', results, getThumbnails(results)
+    return out_vid+'.webm', results, getThumbnails(results), target_poses_coords
     
 
 def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True):