dwpose node
This commit is contained in:
kijai
@@ -0,0 +1 @@
|
||||
*.pyc
|
||||
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .wholebody import Wholebody
|
||||
|
||||
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
class DWposeDetector:
|
||||
"""
|
||||
A pose detect method for image-like data.
|
||||
|
||||
Parameters:
|
||||
model_det: (str) serialized ONNX format model path,
|
||||
such as https://huggingface.co/yzd-v/DWPose/blob/main/yolox_l.onnx
|
||||
model_pose: (str) serialized ONNX format model path,
|
||||
such as https://huggingface.co/yzd-v/DWPose/blob/main/dw-ll_ucoco_384.onnx
|
||||
device: (str) 'cpu' or 'cuda:{device_id}'
|
||||
"""
|
||||
def __init__(self, model_det, model_pose, device='cpu'):
|
||||
self.pose_estimation = Wholebody(model_det=model_det, model_pose=model_pose, device=device)
|
||||
|
||||
def __call__(self, oriImg):
|
||||
oriImg = oriImg.copy()
|
||||
H, W, C = oriImg.shape
|
||||
with torch.no_grad():
|
||||
candidate, score = self.pose_estimation(oriImg)
|
||||
nums, _, locs = candidate.shape
|
||||
candidate[..., 0] /= float(W)
|
||||
candidate[..., 1] /= float(H)
|
||||
body = candidate[:, :18].copy()
|
||||
body = body.reshape(nums * 18, locs)
|
||||
subset = score[:, :18].copy()
|
||||
for i in range(len(subset)):
|
||||
for j in range(len(subset[i])):
|
||||
if subset[i][j] > 0.3:
|
||||
subset[i][j] = int(18 * i + j)
|
||||
else:
|
||||
subset[i][j] = -1
|
||||
|
||||
# un_visible = subset < 0.3
|
||||
# candidate[un_visible] = -1
|
||||
|
||||
# foot = candidate[:, 18:24]
|
||||
|
||||
faces = candidate[:, 24:92]
|
||||
|
||||
hands = candidate[:, 92:113]
|
||||
hands = np.vstack([hands, candidate[:, 113:]])
|
||||
|
||||
faces_score = score[:, 24:92]
|
||||
hands_score = np.vstack([score[:, 92:113], score[:, 113:]])
|
||||
|
||||
bodies = dict(candidate=body, subset=subset, score=score[:, :18])
|
||||
pose = dict(bodies=bodies, hands=hands, hands_score=hands_score, faces=faces, faces_score=faces_score)
|
||||
|
||||
return pose
|
||||
|
||||
# dwpose_detector = DWposeDetector(
|
||||
# model_det="models/DWPose/yolox_l.onnx",
|
||||
# model_pose="models/DWPose/dw-ll_ucoco_384.onnx",
|
||||
# device=device)
|
||||
@@ -0,0 +1,145 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def nms(boxes, scores, nms_thr):
|
||||
"""Single class NMS implemented in Numpy.
|
||||
|
||||
Args:
|
||||
boxes (np.ndarray): shape=(N,4); N is number of boxes
|
||||
scores (np.ndarray): the score of bboxes
|
||||
nms_thr (float): the threshold in NMS
|
||||
|
||||
Returns:
|
||||
List[int]: output bbox ids
|
||||
"""
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = boxes[:, 2]
|
||||
y2 = boxes[:, 3]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= nms_thr)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
def multiclass_nms(boxes, scores, nms_thr, score_thr):
|
||||
"""Multiclass NMS implemented in Numpy. Class-aware version.
|
||||
|
||||
Args:
|
||||
boxes (np.ndarray): shape=(N,4); N is number of boxes
|
||||
scores (np.ndarray): the score of bboxes
|
||||
nms_thr (float): the threshold in NMS
|
||||
score_thr (float): the threshold of cls score
|
||||
|
||||
Returns:
|
||||
np.ndarray: outputs bboxes coordinate
|
||||
"""
|
||||
final_dets = []
|
||||
num_classes = scores.shape[1]
|
||||
for cls_ind in range(num_classes):
|
||||
cls_scores = scores[:, cls_ind]
|
||||
valid_score_mask = cls_scores > score_thr
|
||||
if valid_score_mask.sum() == 0:
|
||||
continue
|
||||
else:
|
||||
valid_scores = cls_scores[valid_score_mask]
|
||||
valid_boxes = boxes[valid_score_mask]
|
||||
keep = nms(valid_boxes, valid_scores, nms_thr)
|
||||
if len(keep) > 0:
|
||||
cls_inds = np.ones((len(keep), 1)) * cls_ind
|
||||
dets = np.concatenate(
|
||||
[valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
|
||||
)
|
||||
final_dets.append(dets)
|
||||
if len(final_dets) == 0:
|
||||
return None
|
||||
return np.concatenate(final_dets, 0)
|
||||
|
||||
def demo_postprocess(outputs, img_size, p6=False):
|
||||
grids = []
|
||||
expanded_strides = []
|
||||
strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
|
||||
|
||||
hsizes = [img_size[0] // stride for stride in strides]
|
||||
wsizes = [img_size[1] // stride for stride in strides]
|
||||
|
||||
for hsize, wsize, stride in zip(hsizes, wsizes, strides):
|
||||
xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
|
||||
grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
|
||||
grids.append(grid)
|
||||
shape = grid.shape[:2]
|
||||
expanded_strides.append(np.full((*shape, 1), stride))
|
||||
|
||||
grids = np.concatenate(grids, 1)
|
||||
expanded_strides = np.concatenate(expanded_strides, 1)
|
||||
outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
|
||||
outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
|
||||
|
||||
return outputs
|
||||
|
||||
def preprocess(img, input_size, swap=(2, 0, 1)):
|
||||
if len(img.shape) == 3:
|
||||
padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
|
||||
else:
|
||||
padded_img = np.ones(input_size, dtype=np.uint8) * 114
|
||||
|
||||
r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
|
||||
resized_img = cv2.resize(
|
||||
img,
|
||||
(int(img.shape[1] * r), int(img.shape[0] * r)),
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
).astype(np.uint8)
|
||||
padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
|
||||
|
||||
padded_img = padded_img.transpose(swap)
|
||||
padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
|
||||
return padded_img, r
|
||||
|
||||
def inference_detector(session, oriImg):
|
||||
"""run human detect
|
||||
"""
|
||||
input_shape = (640,640)
|
||||
img, ratio = preprocess(oriImg, input_shape)
|
||||
|
||||
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
|
||||
output = session.run(None, ort_inputs)
|
||||
predictions = demo_postprocess(output[0], input_shape)[0]
|
||||
|
||||
boxes = predictions[:, :4]
|
||||
scores = predictions[:, 4:5] * predictions[:, 5:]
|
||||
|
||||
boxes_xyxy = np.ones_like(boxes)
|
||||
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
|
||||
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
|
||||
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
|
||||
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
|
||||
boxes_xyxy /= ratio
|
||||
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
|
||||
if dets is not None:
|
||||
final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
|
||||
isscore = final_scores>0.3
|
||||
iscat = final_cls_inds == 0
|
||||
isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
|
||||
final_boxes = final_boxes[isbbox]
|
||||
else:
|
||||
final_boxes = np.array([])
|
||||
|
||||
return final_boxes
|
||||
@@ -0,0 +1,375 @@
|
||||
from typing import List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
def preprocess(
|
||||
img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Do preprocessing for RTMPose model inference.
|
||||
|
||||
Args:
|
||||
img (np.ndarray): Input image in shape.
|
||||
input_size (tuple): Input image size in shape (w, h).
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- resized_img (np.ndarray): Preprocessed image.
|
||||
- center (np.ndarray): Center of image.
|
||||
- scale (np.ndarray): Scale of image.
|
||||
"""
|
||||
# get shape of image
|
||||
img_shape = img.shape[:2]
|
||||
out_img, out_center, out_scale = [], [], []
|
||||
if len(out_bbox) == 0:
|
||||
out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
|
||||
for i in range(len(out_bbox)):
|
||||
x0 = out_bbox[i][0]
|
||||
y0 = out_bbox[i][1]
|
||||
x1 = out_bbox[i][2]
|
||||
y1 = out_bbox[i][3]
|
||||
bbox = np.array([x0, y0, x1, y1])
|
||||
|
||||
# get center and scale
|
||||
center, scale = bbox_xyxy2cs(bbox, padding=1.25)
|
||||
|
||||
# do affine transformation
|
||||
resized_img, scale = top_down_affine(input_size, scale, center, img)
|
||||
|
||||
# normalize image
|
||||
mean = np.array([123.675, 116.28, 103.53])
|
||||
std = np.array([58.395, 57.12, 57.375])
|
||||
resized_img = (resized_img - mean) / std
|
||||
|
||||
out_img.append(resized_img)
|
||||
out_center.append(center)
|
||||
out_scale.append(scale)
|
||||
|
||||
return out_img, out_center, out_scale
|
||||
|
||||
|
||||
def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray:
|
||||
"""Inference RTMPose model.
|
||||
|
||||
Args:
|
||||
sess (ort.InferenceSession): ONNXRuntime session.
|
||||
img (np.ndarray): Input image in shape.
|
||||
|
||||
Returns:
|
||||
outputs (np.ndarray): Output of RTMPose model.
|
||||
"""
|
||||
all_out = []
|
||||
# build input
|
||||
for i in range(len(img)):
|
||||
input = [img[i].transpose(2, 0, 1)]
|
||||
|
||||
# build output
|
||||
sess_input = {sess.get_inputs()[0].name: input}
|
||||
sess_output = []
|
||||
for out in sess.get_outputs():
|
||||
sess_output.append(out.name)
|
||||
|
||||
# run model
|
||||
outputs = sess.run(sess_output, sess_input)
|
||||
all_out.append(outputs)
|
||||
|
||||
return all_out
|
||||
|
||||
|
||||
def postprocess(outputs: List[np.ndarray],
|
||||
model_input_size: Tuple[int, int],
|
||||
center: Tuple[int, int],
|
||||
scale: Tuple[int, int],
|
||||
simcc_split_ratio: float = 2.0
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Postprocess for RTMPose model output.
|
||||
|
||||
Args:
|
||||
outputs (np.ndarray): Output of RTMPose model.
|
||||
model_input_size (tuple): RTMPose model Input image size.
|
||||
center (tuple): Center of bbox in shape (x, y).
|
||||
scale (tuple): Scale of bbox in shape (w, h).
|
||||
simcc_split_ratio (float): Split ratio of simcc.
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- keypoints (np.ndarray): Rescaled keypoints.
|
||||
- scores (np.ndarray): Model predict scores.
|
||||
"""
|
||||
all_key = []
|
||||
all_score = []
|
||||
for i in range(len(outputs)):
|
||||
# use simcc to decode
|
||||
simcc_x, simcc_y = outputs[i]
|
||||
keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
|
||||
|
||||
# rescale keypoints
|
||||
keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
|
||||
all_key.append(keypoints[0])
|
||||
all_score.append(scores[0])
|
||||
|
||||
return np.array(all_key), np.array(all_score)
|
||||
|
||||
|
||||
def bbox_xyxy2cs(bbox: np.ndarray,
|
||||
padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Transform the bbox format from (x,y,w,h) into (center, scale)
|
||||
|
||||
Args:
|
||||
bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
|
||||
as (left, top, right, bottom)
|
||||
padding (float): BBox padding factor that will be multilied to scale.
|
||||
Default: 1.0
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing center and scale.
|
||||
- np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
|
||||
(n, 2)
|
||||
- np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
|
||||
(n, 2)
|
||||
"""
|
||||
# convert single bbox from (4, ) to (1, 4)
|
||||
dim = bbox.ndim
|
||||
if dim == 1:
|
||||
bbox = bbox[None, :]
|
||||
|
||||
# get bbox center and scale
|
||||
x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
|
||||
center = np.hstack([x1 + x2, y1 + y2]) * 0.5
|
||||
scale = np.hstack([x2 - x1, y2 - y1]) * padding
|
||||
|
||||
if dim == 1:
|
||||
center = center[0]
|
||||
scale = scale[0]
|
||||
|
||||
return center, scale
|
||||
|
||||
|
||||
def _fix_aspect_ratio(bbox_scale: np.ndarray,
|
||||
aspect_ratio: float) -> np.ndarray:
|
||||
"""Extend the scale to match the given aspect ratio.
|
||||
|
||||
Args:
|
||||
scale (np.ndarray): The image scale (w, h) in shape (2, )
|
||||
aspect_ratio (float): The ratio of ``w/h``
|
||||
|
||||
Returns:
|
||||
np.ndarray: The reshaped image scale in (2, )
|
||||
"""
|
||||
w, h = np.hsplit(bbox_scale, [1])
|
||||
bbox_scale = np.where(w > h * aspect_ratio,
|
||||
np.hstack([w, w / aspect_ratio]),
|
||||
np.hstack([h * aspect_ratio, h]))
|
||||
return bbox_scale
|
||||
|
||||
|
||||
def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
|
||||
"""Rotate a point by an angle.
|
||||
|
||||
Args:
|
||||
pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
|
||||
angle_rad (float): rotation angle in radian
|
||||
|
||||
Returns:
|
||||
np.ndarray: Rotated point in shape (2, )
|
||||
"""
|
||||
sn, cs = np.sin(angle_rad), np.cos(angle_rad)
|
||||
rot_mat = np.array([[cs, -sn], [sn, cs]])
|
||||
return rot_mat @ pt
|
||||
|
||||
|
||||
def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
|
||||
"""To calculate the affine matrix, three pairs of points are required. This
|
||||
function is used to get the 3rd point, given 2D points a & b.
|
||||
|
||||
The 3rd point is defined by rotating vector `a - b` by 90 degrees
|
||||
anticlockwise, using b as the rotation center.
|
||||
|
||||
Args:
|
||||
a (np.ndarray): The 1st point (x,y) in shape (2, )
|
||||
b (np.ndarray): The 2nd point (x,y) in shape (2, )
|
||||
|
||||
Returns:
|
||||
np.ndarray: The 3rd point.
|
||||
"""
|
||||
direction = a - b
|
||||
c = b + np.r_[-direction[1], direction[0]]
|
||||
return c
|
||||
|
||||
|
||||
def get_warp_matrix(center: np.ndarray,
|
||||
scale: np.ndarray,
|
||||
rot: float,
|
||||
output_size: Tuple[int, int],
|
||||
shift: Tuple[float, float] = (0., 0.),
|
||||
inv: bool = False) -> np.ndarray:
|
||||
"""Calculate the affine transformation matrix that can warp the bbox area
|
||||
in the input image to the output size.
|
||||
|
||||
Args:
|
||||
center (np.ndarray[2, ]): Center of the bounding box (x, y).
|
||||
scale (np.ndarray[2, ]): Scale of the bounding box
|
||||
wrt [width, height].
|
||||
rot (float): Rotation angle (degree).
|
||||
output_size (np.ndarray[2, ] | list(2,)): Size of the
|
||||
destination heatmaps.
|
||||
shift (0-100%): Shift translation ratio wrt the width/height.
|
||||
Default (0., 0.).
|
||||
inv (bool): Option to inverse the affine transform direction.
|
||||
(inv=False: src->dst or inv=True: dst->src)
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 2x3 transformation matrix
|
||||
"""
|
||||
shift = np.array(shift)
|
||||
src_w = scale[0]
|
||||
dst_w = output_size[0]
|
||||
dst_h = output_size[1]
|
||||
|
||||
# compute transformation matrix
|
||||
rot_rad = np.deg2rad(rot)
|
||||
src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
|
||||
dst_dir = np.array([0., dst_w * -0.5])
|
||||
|
||||
# get four corners of the src rectangle in the original image
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = center + scale * shift
|
||||
src[1, :] = center + src_dir + scale * shift
|
||||
src[2, :] = _get_3rd_point(src[0, :], src[1, :])
|
||||
|
||||
# get four corners of the dst rectangle in the input image
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
||||
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
||||
dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
|
||||
|
||||
if inv:
|
||||
warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
||||
else:
|
||||
warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
||||
|
||||
return warp_mat
|
||||
|
||||
|
||||
def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
|
||||
img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Get the bbox image as the model input by affine transform.
|
||||
|
||||
Args:
|
||||
input_size (dict): The input size of the model.
|
||||
bbox_scale (dict): The bbox scale of the img.
|
||||
bbox_center (dict): The bbox center of the img.
|
||||
img (np.ndarray): The original image.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing center and scale.
|
||||
- np.ndarray[float32]: img after affine transform.
|
||||
- np.ndarray[float32]: bbox scale after affine transform.
|
||||
"""
|
||||
w, h = input_size
|
||||
warp_size = (int(w), int(h))
|
||||
|
||||
# reshape bbox to fixed aspect ratio
|
||||
bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
|
||||
|
||||
# get the affine matrix
|
||||
center = bbox_center
|
||||
scale = bbox_scale
|
||||
rot = 0
|
||||
warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
|
||||
|
||||
# do affine transform
|
||||
img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
|
||||
|
||||
return img, bbox_scale
|
||||
|
||||
|
||||
def get_simcc_maximum(simcc_x: np.ndarray,
|
||||
simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Get maximum response location and value from simcc representations.
|
||||
|
||||
Note:
|
||||
instance number: N
|
||||
num_keypoints: K
|
||||
heatmap height: H
|
||||
heatmap width: W
|
||||
|
||||
Args:
|
||||
simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
|
||||
simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- locs (np.ndarray): locations of maximum heatmap responses in shape
|
||||
(K, 2) or (N, K, 2)
|
||||
- vals (np.ndarray): values of maximum heatmap responses in shape
|
||||
(K,) or (N, K)
|
||||
"""
|
||||
N, K, Wx = simcc_x.shape
|
||||
simcc_x = simcc_x.reshape(N * K, -1)
|
||||
simcc_y = simcc_y.reshape(N * K, -1)
|
||||
|
||||
# get maximum value locations
|
||||
x_locs = np.argmax(simcc_x, axis=1)
|
||||
y_locs = np.argmax(simcc_y, axis=1)
|
||||
locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
|
||||
max_val_x = np.amax(simcc_x, axis=1)
|
||||
max_val_y = np.amax(simcc_y, axis=1)
|
||||
|
||||
# get maximum value across x and y axis
|
||||
mask = max_val_x > max_val_y
|
||||
max_val_x[mask] = max_val_y[mask]
|
||||
vals = max_val_x
|
||||
locs[vals <= 0.] = -1
|
||||
|
||||
# reshape
|
||||
locs = locs.reshape(N, K, 2)
|
||||
vals = vals.reshape(N, K)
|
||||
|
||||
return locs, vals
|
||||
|
||||
|
||||
def decode(simcc_x: np.ndarray, simcc_y: np.ndarray,
|
||||
simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""Modulate simcc distribution with Gaussian.
|
||||
|
||||
Args:
|
||||
simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
|
||||
simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
|
||||
simcc_split_ratio (int): The split ratio of simcc.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing center and scale.
|
||||
- np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
|
||||
- np.ndarray[float32]: scores in shape (K,) or (n, K)
|
||||
"""
|
||||
keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
|
||||
keypoints /= simcc_split_ratio
|
||||
|
||||
return keypoints, scores
|
||||
|
||||
|
||||
def inference_pose(session, out_bbox, oriImg):
|
||||
"""run pose detect
|
||||
|
||||
Args:
|
||||
session (ort.InferenceSession): ONNXRuntime session.
|
||||
out_bbox (np.ndarray): bbox list
|
||||
oriImg (np.ndarray): Input image in shape.
|
||||
|
||||
Returns:
|
||||
tuple:
|
||||
- keypoints (np.ndarray): Rescaled keypoints.
|
||||
- scores (np.ndarray): Model predict scores.
|
||||
"""
|
||||
h, w = session.get_inputs()[0].shape[2:]
|
||||
model_input_size = (w, h)
|
||||
# preprocess for rtm-pose model inference.
|
||||
resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
|
||||
# run pose estimation for processed img
|
||||
outputs = inference(session, resized_img)
|
||||
# postprocess for rtm-pose model output.
|
||||
keypoints, scores = postprocess(outputs, model_input_size, center, scale)
|
||||
|
||||
return keypoints, scores
|
||||
@@ -0,0 +1,70 @@
|
||||
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, 11, 14, 15, 16, 17]
|
||||
ref_keypoint_id = [i for i in ref_keypoint_id \
|
||||
if ref_pose['bodies']['score'].shape[0] > 0 and ref_pose['bodies']['score'][0][i] > 0.3]
|
||||
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))
|
||||
|
||||
detected_poses = [dwprocessor(frm) for frm in vr.get_batch(list(range(0, len(vr), sample_stride))).asnumpy()]
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,136 @@
|
||||
import math
|
||||
import numpy as np
|
||||
import matplotlib
|
||||
import cv2
|
||||
|
||||
|
||||
eps = 0.01
|
||||
|
||||
def alpha_blend_color(color, alpha):
|
||||
"""blend color according to point conf
|
||||
"""
|
||||
return [int(c * alpha) for c in color]
|
||||
|
||||
def draw_bodypose(canvas, candidate, subset, score):
|
||||
H, W, C = canvas.shape
|
||||
candidate = np.array(candidate)
|
||||
subset = np.array(subset)
|
||||
|
||||
stickwidth = 4
|
||||
|
||||
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
|
||||
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
|
||||
[1, 16], [16, 18], [3, 17], [6, 18]]
|
||||
|
||||
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
|
||||
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
|
||||
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
|
||||
|
||||
for i in range(17):
|
||||
for n in range(len(subset)):
|
||||
index = subset[n][np.array(limbSeq[i]) - 1]
|
||||
conf = score[n][np.array(limbSeq[i]) - 1]
|
||||
if conf[0] < 0.3 or conf[1] < 0.3:
|
||||
continue
|
||||
Y = candidate[index.astype(int), 0] * float(W)
|
||||
X = candidate[index.astype(int), 1] * float(H)
|
||||
mX = np.mean(X)
|
||||
mY = np.mean(Y)
|
||||
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
|
||||
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
||||
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
|
||||
cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(colors[i], conf[0] * conf[1]))
|
||||
|
||||
canvas = (canvas * 0.6).astype(np.uint8)
|
||||
|
||||
for i in range(18):
|
||||
for n in range(len(subset)):
|
||||
index = int(subset[n][i])
|
||||
if index == -1:
|
||||
continue
|
||||
x, y = candidate[index][0:2]
|
||||
conf = score[n][i]
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
cv2.circle(canvas, (int(x), int(y)), 4, alpha_blend_color(colors[i], conf), thickness=-1)
|
||||
|
||||
return canvas
|
||||
|
||||
def draw_handpose(canvas, all_hand_peaks, all_hand_scores):
|
||||
H, W, C = canvas.shape
|
||||
|
||||
edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
|
||||
[10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
|
||||
|
||||
for peaks, scores in zip(all_hand_peaks, all_hand_scores):
|
||||
|
||||
for ie, e in enumerate(edges):
|
||||
x1, y1 = peaks[e[0]]
|
||||
x2, y2 = peaks[e[1]]
|
||||
x1 = int(x1 * W)
|
||||
y1 = int(y1 * H)
|
||||
x2 = int(x2 * W)
|
||||
y2 = int(y2 * H)
|
||||
score = int(scores[e[0]] * scores[e[1]] * 255)
|
||||
if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
|
||||
cv2.line(canvas, (x1, y1), (x2, y2),
|
||||
matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * score, thickness=2)
|
||||
|
||||
for i, keyponit in enumerate(peaks):
|
||||
x, y = keyponit
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
score = int(scores[i] * 255)
|
||||
if x > eps and y > eps:
|
||||
cv2.circle(canvas, (x, y), 4, (0, 0, score), thickness=-1)
|
||||
return canvas
|
||||
|
||||
def draw_facepose(canvas, all_lmks, all_scores):
|
||||
H, W, C = canvas.shape
|
||||
for lmks, scores in zip(all_lmks, all_scores):
|
||||
for lmk, score in zip(lmks, scores):
|
||||
x, y = lmk
|
||||
x = int(x * W)
|
||||
y = int(y * H)
|
||||
conf = int(score * 255)
|
||||
if x > eps and y > eps:
|
||||
cv2.circle(canvas, (x, y), 3, (conf, conf, conf), thickness=-1)
|
||||
return canvas
|
||||
|
||||
def draw_pose(pose, H, W, include_body, include_hand, include_face, ref_w=2160):
|
||||
"""vis dwpose outputs
|
||||
|
||||
Args:
|
||||
pose (List): DWposeDetector outputs in dwpose_detector.py
|
||||
H (int): height
|
||||
W (int): width
|
||||
ref_w (int, optional) Defaults to 2160.
|
||||
|
||||
Returns:
|
||||
np.ndarray: image pixel value in RGB mode
|
||||
"""
|
||||
bodies = pose['bodies']
|
||||
faces = pose['faces']
|
||||
hands = pose['hands']
|
||||
candidate = bodies['candidate']
|
||||
subset = bodies['subset']
|
||||
|
||||
sz = min(H, W)
|
||||
sr = (ref_w / sz) if sz != ref_w else 1
|
||||
|
||||
########################################## create zero canvas ##################################################
|
||||
canvas = np.zeros(shape=(int(H*sr), int(W*sr), 3), dtype=np.uint8)
|
||||
|
||||
########################################### draw body pose #####################################################
|
||||
if include_body:
|
||||
canvas = draw_bodypose(canvas, candidate, subset, score=bodies['score'])
|
||||
|
||||
########################################### draw hand pose #####################################################
|
||||
if include_hand:
|
||||
canvas = draw_handpose(canvas, hands, pose['hands_score'])
|
||||
|
||||
########################################### draw face pose #####################################################
|
||||
if include_face:
|
||||
canvas = draw_facepose(canvas, faces, pose['faces_score'])
|
||||
|
||||
return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
|
||||
@@ -0,0 +1,57 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
from .onnxdet import inference_detector
|
||||
from .onnxpose import inference_pose
|
||||
|
||||
|
||||
class Wholebody:
|
||||
"""detect human pose by dwpose
|
||||
"""
|
||||
def __init__(self, model_det, model_pose, device="cpu"):
|
||||
providers = ['CPUExecutionProvider'] if device == 'cpu' else ['CUDAExecutionProvider']
|
||||
provider_options = None if device == 'cpu' else [{'device_id': 0}]
|
||||
|
||||
self.session_det = ort.InferenceSession(
|
||||
path_or_bytes=model_det, providers=providers, provider_options=provider_options
|
||||
)
|
||||
self.session_pose = ort.InferenceSession(
|
||||
path_or_bytes=model_pose, providers=providers, provider_options=provider_options
|
||||
)
|
||||
|
||||
def __call__(self, oriImg):
|
||||
"""call to process dwpose-detect
|
||||
|
||||
Args:
|
||||
oriImg (np.ndarray): detected image
|
||||
|
||||
"""
|
||||
det_result = inference_detector(self.session_det, oriImg)
|
||||
keypoints, scores = inference_pose(self.session_pose, det_result, oriImg)
|
||||
|
||||
keypoints_info = np.concatenate(
|
||||
(keypoints, scores[..., None]), axis=-1)
|
||||
# compute neck joint
|
||||
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
|
||||
# neck score when visualizing pred
|
||||
neck[:, 2:4] = np.logical_and(
|
||||
keypoints_info[:, 5, 2:4] > 0.3,
|
||||
keypoints_info[:, 6, 2:4] > 0.3).astype(int)
|
||||
new_keypoints_info = np.insert(
|
||||
keypoints_info, 17, neck, axis=1)
|
||||
mmpose_idx = [
|
||||
17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
|
||||
]
|
||||
openpose_idx = [
|
||||
1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
|
||||
]
|
||||
new_keypoints_info[:, openpose_idx] = \
|
||||
new_keypoints_info[:, mmpose_idx]
|
||||
keypoints_info = new_keypoints_info
|
||||
|
||||
keypoints, scores = keypoints_info[
|
||||
..., :2], keypoints_info[..., 2]
|
||||
|
||||
return keypoints, scores
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ from omegaconf import OmegaConf
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import sys
|
||||
import numpy as np
|
||||
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(script_directory)
|
||||
@@ -27,7 +28,6 @@ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
|
||||
from mimicmotion.modules.unet import UNetSpatioTemporalConditionModel
|
||||
from mimicmotion.modules.pose_net import PoseNet
|
||||
from mimicmotion.pipelines.pipeline_mimicmotion import MimicMotionPipeline
|
||||
|
||||
class MimicMotionModel(torch.nn.Module):
|
||||
def __init__(self, base_model_path):
|
||||
@@ -182,13 +182,114 @@ class MimicMotionSampler:
|
||||
|
||||
return frames,
|
||||
|
||||
class MimicMotionGetPoses:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"ref_image": ("IMAGE",),
|
||||
"pose_images": ("IMAGE",),
|
||||
"include_body": ("BOOLEAN", {"default": True}),
|
||||
"include_hand": ("BOOLEAN", {"default": True}),
|
||||
"include_face": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "MimicMotionWrapper"
|
||||
|
||||
def process(self, ref_image, pose_images, include_body, include_hand, include_face):
|
||||
device = mm.get_torch_device()
|
||||
from mimicmotion.dwpose.util import draw_pose
|
||||
from mimicmotion.dwpose.dwpose_detector import DWposeDetector
|
||||
|
||||
yolo_model = "yolox_l.onnx"
|
||||
dw_pose_model = "dw-ll_ucoco_384.onnx"
|
||||
model_base_path = os.path.join(script_directory, "models", "DWPose")
|
||||
|
||||
model_det=os.path.join(model_base_path, yolo_model)
|
||||
model_pose=os.path.join(model_base_path, dw_pose_model)
|
||||
|
||||
if not os.path.exists(model_det):
|
||||
print(f"Downloading yolo model to: {model_base_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id="yzd-v/DWPose",
|
||||
allow_patterns=[f"*{yolo_model}*"],
|
||||
local_dir=model_base_path,
|
||||
local_dir_use_symlinks=False)
|
||||
|
||||
if not os.path.exists(model_pose):
|
||||
print(f"Downloading dwpose model to: {model_base_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id="yzd-v/DWPose",
|
||||
allow_patterns=[f"*{dw_pose_model}*"],
|
||||
local_dir=model_base_path,
|
||||
local_dir_use_symlinks=False)
|
||||
|
||||
dwprocessor = DWposeDetector(
|
||||
model_det=os.path.join(model_base_path, "yolox_l.onnx"),
|
||||
model_pose=os.path.join(model_base_path, "dw-ll_ucoco_384.onnx"),
|
||||
device=device)
|
||||
|
||||
ref_image = ref_image.squeeze(0).cpu().numpy() * 255
|
||||
|
||||
# select ref-keypoint from reference pose for pose rescale
|
||||
ref_pose = dwprocessor(ref_image)
|
||||
ref_keypoint_id = [0, 1, 2, 5, 8, 11, 14, 15, 16, 17]
|
||||
ref_keypoint_id = [i for i in ref_keypoint_id \
|
||||
if ref_pose['bodies']['score'].shape[0] > 0 and ref_pose['bodies']['score'][0][i] > 0.3]
|
||||
ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id]
|
||||
|
||||
|
||||
height, width, _ = ref_image.shape
|
||||
pose_images_np = pose_images.cpu().numpy() * 255
|
||||
|
||||
# read input video
|
||||
detected_poses_np_list = []
|
||||
for img_np in pose_images_np:
|
||||
detected_poses_np_list.append(dwprocessor(img_np))
|
||||
|
||||
detected_bodies = np.stack(
|
||||
[p['bodies']['candidate'] for p in detected_poses_np_list 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, _ = pose_images_np[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_np_list:
|
||||
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, include_body=include_body, include_hand=include_hand, include_face=include_face)
|
||||
output_pose.append(np.array(im))
|
||||
|
||||
output_pose_tensors = [torch.tensor(np.array(im)) for im in output_pose]
|
||||
output_tensor = torch.stack(output_pose_tensors) / 255
|
||||
|
||||
ref_pose_img = draw_pose(ref_pose, height, width, include_body=include_body, include_hand=include_hand, include_face=include_face)
|
||||
ref_pose_tensor = torch.tensor(np.array(ref_pose_img)) / 255
|
||||
output_tensor = torch.cat((ref_pose_tensor.unsqueeze(0), output_tensor))
|
||||
output_tensor = output_tensor.permute(0, 2, 3, 1).cpu().float()
|
||||
|
||||
return output_tensor,
|
||||
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DownloadAndLoadMimicMotionModel": DownloadAndLoadMimicMotionModel,
|
||||
"MimicMotionSampler": MimicMotionSampler,
|
||||
"MimicMotionGetPoses": MimicMotionGetPoses
|
||||
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DownloadAndLoadMimicMotionModel": "DownloadAndLoadMimicMotionModel",
|
||||
"MimicMotionSampler": "MimicMotionSampler",
|
||||
"MimicMotionGetPoses": "MimicMotionGetPoses"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user