in src/controlnet_aux/dwpose/wholebody.py [0:0]
def __call__(self, oriImg):
# predict bbox
det_result = inference_detector(self.detector, oriImg)
pred_instance = det_result.pred_instances.cpu().numpy()
bboxes = np.concatenate(
(pred_instance.bboxes, pred_instance.scores[:, None]), axis=1)
bboxes = bboxes[np.logical_and(pred_instance.labels == 0,
pred_instance.scores > 0.5)]
# set NMS threshold
bboxes = bboxes[nms(bboxes, 0.7), :4]
# predict keypoints
if len(bboxes) == 0:
pose_results = inference_topdown(self.pose_estimator, oriImg)
else:
pose_results = inference_topdown(self.pose_estimator, oriImg, bboxes)
preds = merge_data_samples(pose_results)
preds = preds.pred_instances
# preds = pose_results[0].pred_instances
keypoints = preds.get('transformed_keypoints',
preds.keypoints)
if 'keypoint_scores' in preds:
scores = preds.keypoint_scores
else:
scores = np.ones(keypoints.shape[:-1])
if 'keypoints_visible' in preds:
visible = preds.keypoints_visible
else:
visible = np.ones(keypoints.shape[:-1])
keypoints_info = np.concatenate(
(keypoints, scores[..., None], visible[..., 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, visible = keypoints_info[
..., :2], keypoints_info[..., 2], keypoints_info[..., 3]
return keypoints, scores