crosstyan/RapidPoseTriangulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Moved detector pre/post-processing into onnx graph.

Browse Source
This commit is contained in:
Daniel
2024-12-03 11:52:55 +01:00
parent 36781e616b
commit 742d2386c7
5 changed files with 213 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
scripts/utils_2d_pose_ort.py
View File

@ -31,6 +31,8 @@ class BaseModel(ABC):
self.input_name = self.session.get_inputs()[0].name
self.input_shape = self.session.get_inputs()[0].shape
if "batch_size" in self.input_shape:
self.input_shape = [1, 500, 500, 3]
input_type = self.session.get_inputs()[0].type
if input_type == "tensor(float16)":
@ -72,116 +74,25 @@ class RTMDet(BaseModel):
self,
model_path: str,
conf_threshold: float,
iou_threshold: float,
warmup: int = 30,
):
super(RTMDet, self).__init__(model_path, warmup)
self.conf_threshold = conf_threshold
self.iou_threshold = iou_threshold
self.dx = 0
self.dy = 0
self.scale = 0
def letterbox(self, img: np.ndarray, target_size: List[int], fill_value: int = 128):
h, w = img.shape[:2]
tw, th = target_size
scale = min(tw / w, th / h)
nw, nh = int(w * scale), int(h * scale)
dx, dy = (tw - nw) // 2, (th - nh) // 2
canvas = np.full((th, tw, img.shape[2]), fill_value, dtype=img.dtype)
canvas[dy : dy + nh, dx : dx + nw, :] = cv2.resize(
img, (nw, nh), interpolation=cv2.INTER_LINEAR
)
return canvas, dx, dy, scale
def nms_optimized(
self, boxes: np.ndarray, iou_threshold: float, conf_threshold: float
):
"""
Perform Non-Maximum Suppression (NMS) on bounding boxes for a single class.
"""
# Filter out boxes with low confidence scores
scores = boxes[:, 4]
keep = scores > conf_threshold
boxes = boxes[keep]
scores = scores[keep]
if boxes.shape[0] == 0:
return np.empty((0, 5), dtype=boxes.dtype)
# Compute the area of the bounding boxes
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
# Sort the boxes by scores in descending order
order = scores.argsort()[::-1]
keep_indices = []
while order.size > 0:
i = order[0]
keep_indices.append(i)
# Compute IoU of the current box with the rest
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:]])
# Compute width and height of the overlapping area
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
# Compute the area of the intersection
inter = w * h
# Compute the IoU
iou = inter / (areas[i] + areas[order[1:]] - inter)
# Keep boxes with IoU less than the threshold
inds = np.where(iou <= iou_threshold)[0]
# Update the order array
order = order[inds + 1]
# Return the boxes that are kept
return boxes[keep_indices]
def preprocess(self, image: np.ndarray):
th, tw = self.input_shape[1:3]
image, self.dx, self.dy, self.scale = self.letterbox(
image, (tw, th), fill_value=114
)
tensor = np.asarray(image).astype(self.input_type, copy=False)
tensor = np.expand_dims(tensor, axis=0)
return tensor
def postprocess(self, tensor: List[np.ndarray]):
boxes = np.squeeze(tensor[0], axis=0)
classes = np.expand_dims(np.squeeze(tensor[1], axis=0), axis=-1)
boxes = np.concatenate([boxes, classes], axis=-1)
boxes = np.squeeze(tensor[1], axis=0)
classes = np.squeeze(tensor[0], axis=0)
boxes = self.nms_optimized(boxes, self.iou_threshold, self.conf_threshold)
human_class = classes[:] == 0
boxes = boxes[human_class]
if boxes.shape[0] == 0:
return boxes
human_class = boxes[..., -1] == 0
boxes = boxes[human_class][..., :4]
boxes[:, 0] -= self.dx
boxes[:, 2] -= self.dx
boxes[:, 1] -= self.dy
boxes[:, 3] -= self.dy
boxes = np.clip(boxes, a_min=0, a_max=None)
boxes[:, :4] /= self.scale
keep = boxes[:, 4] > self.conf_threshold
boxes = boxes[keep]
return boxes
@ -201,7 +112,7 @@ class RTMPose(BaseModel):
target_size: List[int],
padding_scale: float = 1.25,
):
start_x, start_y, end_x, end_y = box
start_x, start_y, end_x, end_y = box[0:4]
target_w, target_h = target_size
# Calculate original bounding box width and height
@ -305,8 +216,7 @@ class TopDown:
self,
det_model_path,
pose_model_path,
conf_threshold=0.6,
iou_threshold=0.6,
box_conf_threshold=0.6,
warmup=30,
):
if (not det_model_path.endswith(".onnx")) or (
@ -314,7 +224,7 @@ class TopDown:
):
raise ValueError("Only ONNX models are supported.")
self.det_model = RTMDet(det_model_path, conf_threshold, iou_threshold, warmup)
self.det_model = RTMDet(det_model_path, box_conf_threshold, warmup)
self.pose_model = RTMPose(pose_model_path, warmup)
def predict(self, image):
@ -337,8 +247,7 @@ def load_model():
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
# "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_extra-steps.onnx",
"/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
conf_threshold=0.3,
iou_threshold=0.3,
box_conf_threshold=0.3,
warmup=30,
)