129 lines
4.6 KiB
Python
129 lines
4.6 KiB
Python
from typing import List, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ..base import BaseTool
|
|
from ..object_detection.post_processings import multiclass_nms
|
|
from .post_processings import convert_coco_to_openpose
|
|
|
|
|
|
class RTMO(BaseTool):
|
|
|
|
def __init__(self,
|
|
onnx_model: str,
|
|
model_input_size: tuple = (640, 640),
|
|
mean: tuple = None,
|
|
std: tuple = None,
|
|
nms_thr: float = 0.45,
|
|
score_thr: float = 0.7,
|
|
to_openpose: bool = False,
|
|
backend: str = 'onnxruntime',
|
|
device: str = 'cpu'):
|
|
super().__init__(onnx_model, model_input_size, mean, std, backend,
|
|
device)
|
|
self.to_openpose = to_openpose
|
|
self.nms_thr = nms_thr
|
|
self.score_thr = score_thr
|
|
|
|
def __call__(self,
|
|
image: np.ndarray,
|
|
nms_thr: float = None,
|
|
score_thr: float = None):
|
|
nms_thr = nms_thr if nms_thr is not None else self.nms_thr
|
|
score_thr = score_thr if score_thr is not None else self.score_thr
|
|
|
|
image, ratio = self.preprocess(image)
|
|
outputs = self.inference(image)
|
|
|
|
keypoints, scores = self.postprocess(outputs, ratio, nms_thr,
|
|
score_thr)
|
|
|
|
if self.to_openpose:
|
|
keypoints, scores = convert_coco_to_openpose(keypoints, scores)
|
|
|
|
return keypoints, scores
|
|
|
|
def preprocess(self, img: np.ndarray):
|
|
"""Do preprocessing for RTMPose model inference.
|
|
|
|
Args:
|
|
img (np.ndarray): Input image in shape.
|
|
|
|
Returns:
|
|
tuple:
|
|
- padded_img (np.ndarray): Preprocessed image.
|
|
- ratio (float): Scale factor applied to the image.
|
|
"""
|
|
if img.shape[:2] == tuple(self.model_input_size[:2]):
|
|
padded_img = img.copy()
|
|
ratio = 1.
|
|
else:
|
|
if len(img.shape) == 3:
|
|
padded_img = np.ones(
|
|
(self.model_input_size[0], self.model_input_size[1], 3),
|
|
dtype=np.uint8) * 114
|
|
else:
|
|
padded_img = np.ones(self.model_input_size,
|
|
dtype=np.uint8) * 114
|
|
|
|
ratio = min(self.model_input_size[0] / img.shape[0],
|
|
self.model_input_size[1] / img.shape[1])
|
|
resized_img = cv2.resize(
|
|
img,
|
|
(int(img.shape[1] * ratio), int(img.shape[0] * ratio)),
|
|
interpolation=cv2.INTER_LINEAR,
|
|
).astype(np.uint8)
|
|
padded_shape = (int(img.shape[0] * ratio),
|
|
int(img.shape[1] * ratio))
|
|
padded_img[:padded_shape[0], :padded_shape[1]] = resized_img
|
|
|
|
# normalize image
|
|
if self.mean is not None:
|
|
self.mean = np.array(self.mean)
|
|
self.std = np.array(self.std)
|
|
padded_img = (padded_img - self.mean) / self.std
|
|
|
|
return padded_img, ratio
|
|
|
|
def postprocess(
|
|
self,
|
|
outputs: List[np.ndarray],
|
|
ratio: float = 1.,
|
|
nms_thr: float = None,
|
|
score_thr: float = None,
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Do postprocessing for RTMO model inference.
|
|
|
|
Args:
|
|
outputs (List[np.ndarray]): Outputs of RTMO model.
|
|
ratio (float): Ratio of preprocessing.
|
|
|
|
Returns:
|
|
tuple:
|
|
- final_boxes (np.ndarray): Final bounding boxes.
|
|
- final_scores (np.ndarray): Final scores.
|
|
"""
|
|
det_outputs, pose_outputs = outputs
|
|
|
|
# onnx contains nms module (?)
|
|
final_boxes, final_scores = (det_outputs[0, :, :4], det_outputs[0, :,
|
|
4])
|
|
final_boxes /= ratio
|
|
keypoints, scores = pose_outputs[0, :, :, :2], pose_outputs[0, :, :, 2]
|
|
keypoints = keypoints / ratio
|
|
|
|
# apply nms
|
|
dets, keep = multiclass_nms(final_boxes,
|
|
final_scores[:, np.newaxis],
|
|
nms_thr=nms_thr,
|
|
score_thr=score_thr)
|
|
if keep is not None:
|
|
keypoints = keypoints[keep]
|
|
scores = scores[keep]
|
|
else:
|
|
keypoints = np.expand_dims(np.zeros_like(keypoints[0]), axis=0)
|
|
scores = np.expand_dims(np.zeros_like(scores[0]), axis=0)
|
|
|
|
return keypoints, scores
|