Spaces:
Runtime error
Runtime error
File size: 2,618 Bytes
e394497 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from tqdm import tqdm
import decord
import numpy as np
from .util import draw_pose
from .dwpose_detector import dwpose_detector as dwprocessor
def get_video_pose(
video_path: str,
ref_image: np.ndarray,
sample_stride: int=1):
"""preprocess ref image pose and video pose
Args:
video_path (str): video pose path
ref_image (np.ndarray): reference image
sample_stride (int, optional): Defaults to 1.
Returns:
np.ndarray: sequence of video pose
"""
# select ref-keypoint from reference pose for pose rescale
ref_pose = dwprocessor(ref_image)
ref_keypoint_id = [0, 1, 2, 5, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
ref_keypoint_id = [i for i in ref_keypoint_id \
if len(ref_pose['bodies']['subset']) > 0 and ref_pose['bodies']['subset'][0][i] >= .0]
ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id]
height, width, _ = ref_image.shape
# read input video
vr = decord.VideoReader(video_path, ctx=decord.cpu(0))
sample_stride *= max(1, int(vr.get_avg_fps() / 24))
frames = vr.get_batch(list(range(0, len(vr), sample_stride))).asnumpy()
detected_poses = [dwprocessor(frm) for frm in tqdm(frames, desc="DWPose")]
dwprocessor.release_memory()
detected_bodies = np.stack(
[p['bodies']['candidate'] for p in detected_poses if p['bodies']['candidate'].shape[0] == 18])[:,
ref_keypoint_id]
# compute linear-rescale params
ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1)
fh, fw, _ = vr[0].shape
ax = ay / (fh / fw / height * width)
bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax)
a = np.array([ax, ay])
b = np.array([bx, by])
output_pose = []
# pose rescale
for detected_pose in detected_poses:
detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b
detected_pose['faces'] = detected_pose['faces'] * a + b
detected_pose['hands'] = detected_pose['hands'] * a + b
im = draw_pose(detected_pose, height, width)
output_pose.append(np.array(im))
return np.stack(output_pose)
def get_image_pose(ref_image):
"""process image pose
Args:
ref_image (np.ndarray): reference image pixel value
Returns:
np.ndarray: pose visual image in RGB-mode
"""
height, width, _ = ref_image.shape
ref_pose = dwprocessor(ref_image)
pose_img = draw_pose(ref_pose, height, width)
return np.array(pose_img)
|