MCPcopy
hub / github.com/triple-mu/YOLOv8-TensorRT / pose_postprocess

Function pose_postprocess

models/utils.py:169–194  ·  view source on GitHub ↗
(
    data: tuple | ndarray, conf_thres: float = 0.25, iou_thres: float = 0.65
)

Source from the content-addressed store, hash-verified

167
168
169def pose_postprocess(
170 data: tuple | ndarray, conf_thres: float = 0.25, iou_thres: float = 0.65
171) -> tuple[ndarray, ndarray, ndarray, ndarray]:
172 if isinstance(data, tuple):
173 assert len(data) == 1
174 data = data[0]
175 outputs = np.transpose(data[0], (1, 0))
176 num_cls = outputs.shape[-1] - 4 - 51 # 51 = 17 keypoints x 3; usually 1 (person)
177 bboxes, cls, kpts = np.split(outputs, [4, 4 + num_cls], 1)
178 scores = cls.max(-1)
179 labels = cls.argmax(-1)
180 idx = scores > conf_thres
181 if not idx.any(): # no bounding boxes were created
182 return (
183 np.empty((0, 4), dtype=np.float32),
184 np.empty((0,), dtype=np.float32),
185 np.empty((0, 0, 0), dtype=np.float32),
186 np.empty((0,), dtype=np.int32),
187 )
188 bboxes, scores, kpts, labels = bboxes[idx], scores[idx], kpts[idx], labels[idx]
189 xycenter, wh = np.split(bboxes, [2], -1)
190 cvbboxes = np.concatenate([xycenter - 0.5 * wh, wh], -1)
191 nms = cv2.dnn.NMSBoxes(cvbboxes, scores, conf_thres, iou_thres)
192 cvbboxes, scores, kpts, labels = cvbboxes[nms], scores[nms], kpts[nms], labels[nms]
193 cvbboxes[:, 2:] += cvbboxes[:, :2]
194 return cvbboxes, scores, kpts.reshape(nms.shape[0], -1, 3), labels
195
196
197def obb_postprocess(

Calls

no outgoing calls