diff --git a/inference/core/models/utils/keypoints.py b/inference/core/models/utils/keypoints.py index d6d776d463..ffa85ba4b9 100644 --- a/inference/core/models/utils/keypoints.py +++ b/inference/core/models/utils/keypoints.py @@ -22,20 +22,36 @@ def model_keypoints_to_response( if keypoints_metadata is None: raise ModelArtefactError("Keypoints metadata not available.") keypoint_id2name = keypoints_metadata[predicted_object_class_id] + num_kpt = min( + len(keypoints) // 3, len(keypoint_id2name) + ) # pre-calculate loop length for efficiency + results = [] - for keypoint_id in range(len(keypoints) // 3): - if keypoint_id >= len(keypoint_id2name): - # Ultralytics only supports single class keypoint detection, so points might be padded with zeros - break - confidence = keypoints[3 * keypoint_id + 2] - if confidence < keypoint_confidence_threshold: + + # Hoist allocations out of the loop for performance + class_kw = {"class": None} + + # Local bindings for performance + kpt = keypoints + kpt_id2n = keypoint_id2name + kpt_thr = keypoint_confidence_threshold + Keypoint_cls = Keypoint + + # Loop unrolling reduces index calculations + for keypoint_id in range(num_kpt): + idx = 3 * keypoint_id + confidence = kpt[idx + 2] + if confidence < kpt_thr: continue - keypoint = Keypoint( - x=keypoints[3 * keypoint_id], - y=keypoints[3 * keypoint_id + 1], - confidence=confidence, - class_id=keypoint_id, - **{"class": keypoint_id2name[keypoint_id]}, + class_kw["class"] = kpt_id2n[keypoint_id] + results.append( + Keypoint_cls( + x=kpt[idx], + y=kpt[idx + 1], + confidence=confidence, + class_id=keypoint_id, + **class_kw, + ) ) - results.append(keypoint) + return results