data_juicer.ops.common.dwpose_func module¶
- class data_juicer.ops.common.dwpose_func.Wholebody(onnx_det, onnx_pose, device)[source]¶
Bases:
object
- data_juicer.ops.common.dwpose_func.nms(boxes, scores, nms_thr)[source]¶
Single class NMS implemented in Numpy.
- data_juicer.ops.common.dwpose_func.multiclass_nms(boxes, scores, nms_thr, score_thr)[source]¶
Multiclass NMS implemented in Numpy. Class-aware version.
- data_juicer.ops.common.dwpose_func.preprocess_pose(img: ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)) Tuple[ndarray, ndarray, ndarray][source]¶
Do preprocessing for RTMPose model inference.
- Parameters:
img (np.ndarray) – Input image in shape.
input_size (tuple) – Input image size in shape (w, h).
- Returns:
resized_img (np.ndarray): Preprocessed image.
center (np.ndarray): Center of image.
scale (np.ndarray): Scale of image.
- Return type:
tuple
- data_juicer.ops.common.dwpose_func.inference(sess: InferenceSession, img: ndarray) ndarray[source]¶
Inference RTMPose model.
- Parameters:
sess (onnxruntime.InferenceSession) – ONNXRuntime session.
img (np.ndarray) – Input image in shape.
- Returns:
Output of RTMPose model.
- Return type:
outputs (np.ndarray)
- data_juicer.ops.common.dwpose_func.postprocess(outputs: List[ndarray], model_input_size: Tuple[int, int], center: Tuple[int, int], scale: Tuple[int, int], simcc_split_ratio: float = 2.0) Tuple[ndarray, ndarray][source]¶
Postprocess for RTMPose model output.
- Parameters:
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:
keypoints (np.ndarray): Rescaled keypoints.
scores (np.ndarray): Model predict scores.
- Return type:
tuple
- data_juicer.ops.common.dwpose_func.bbox_xyxy2cs(bbox: ndarray, padding: float = 1.0) Tuple[ndarray, ndarray][source]¶
Transform the bbox format from (x,y,w,h) into (center, scale)
- Parameters:
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:
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)
- Return type:
tuple
- data_juicer.ops.common.dwpose_func.get_warp_matrix(center: ndarray, scale: ndarray, rot: float, output_size: Tuple[int, int], shift: Tuple[float, float] = (0.0, 0.0), inv: bool = False) ndarray[source]¶
Calculate the affine transformation matrix that can warp the bbox area in the input image to the output size.
- Parameters:
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:
A 2x3 transformation matrix
- Return type:
np.ndarray
- data_juicer.ops.common.dwpose_func.top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict, img: ndarray) Tuple[ndarray, ndarray][source]¶
Get the bbox image as the model input by affine transform.
- Parameters:
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:
A tuple containing center and scale. - np.ndarray[float32]: img after affine transform. - np.ndarray[float32]: bbox scale after affine transform.
- Return type:
tuple
- data_juicer.ops.common.dwpose_func.get_simcc_maximum(simcc_x: ndarray, simcc_y: ndarray) Tuple[ndarray, ndarray][source]¶
Get maximum response location and value from simcc representations.
Note
instance number: N num_keypoints: K heatmap height: H heatmap width: W
- Parameters:
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:
- 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)
- Return type:
tuple
- data_juicer.ops.common.dwpose_func.decode(simcc_x: ndarray, simcc_y: ndarray, simcc_split_ratio) Tuple[ndarray, ndarray][source]¶
Modulate simcc distribution with Gaussian.
- Parameters:
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:
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)
- Return type:
tuple