148 lines
5.5 KiB
Python
148 lines
5.5 KiB
Python
from typing import List, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ..base import BaseTool
|
|
from .post_processings import multiclass_nms
|
|
|
|
|
|
class RTMDet(BaseTool):
|
|
|
|
def __init__(self,
|
|
onnx_model: str,
|
|
model_input_size: tuple = (640, 640),
|
|
mode: str = 'human',
|
|
mean: tuple = (103.5300, 116.2800, 123.6750),
|
|
std: tuple = (57.3750, 57.1200, 58.3950),
|
|
backend: str = 'onnxruntime',
|
|
device: str = 'cpu'):
|
|
super().__init__(onnx_model,
|
|
model_input_size,
|
|
mean=mean,
|
|
std=std,
|
|
backend=backend,
|
|
device=device)
|
|
self.mode = mode
|
|
|
|
def __call__(self, image: np.ndarray):
|
|
image, ratio = self.preprocess(image)
|
|
outputs = self.inference(image)[0]
|
|
results = self.postprocess(outputs, ratio)
|
|
return results
|
|
|
|
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.,
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Do postprocessing for RTMDet model inference.
|
|
|
|
Args:
|
|
outputs (List[np.ndarray]): Outputs of RTMDet model.
|
|
ratio (float): Ratio of preprocessing.
|
|
|
|
Returns:
|
|
tuple:
|
|
- final_boxes (np.ndarray): Final bounding boxes.
|
|
- final_scores (np.ndarray): Final scores.
|
|
"""
|
|
|
|
if outputs.shape[-1] == 4:
|
|
# onnx without nms module
|
|
|
|
grids = []
|
|
expanded_strides = []
|
|
strides = [8, 16, 32]
|
|
|
|
hsizes = [self.model_input_size[0] // stride for stride in strides]
|
|
wsizes = [self.model_input_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
|
|
|
|
predictions = outputs[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, keep = multiclass_nms(boxes_xyxy,
|
|
scores,
|
|
nms_thr=self.nms_thr,
|
|
score_thr=self.score_thr)
|
|
if dets is not None:
|
|
pack_dets = (dets[:, :4], dets[:, 4], dets[:, 5])
|
|
final_boxes, final_scores, final_cls_inds = pack_dets
|
|
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]
|
|
|
|
elif outputs.shape[-1] == 5:
|
|
# onnx contains nms module
|
|
|
|
pack_dets = (outputs[0, :, :4], outputs[0, :, 4])
|
|
final_boxes, final_scores = pack_dets
|
|
final_boxes /= ratio
|
|
isscore = final_scores > 0.3
|
|
isbbox = [i for i in isscore]
|
|
final_boxes = final_boxes[isbbox]
|
|
|
|
return final_boxes
|