Add a basic from_supervision function to KeyPoints

This commit is contained in:
Linas Kondrackis 2024-04-24 16:37:41 +03:00
parent 2c90a1ef1e
commit 48809c1332
1 changed files with 31 additions and 0 deletions

View File

@ -110,6 +110,37 @@ class KeyPoints:
]
)
@classmethod
def from_inference(cls, inference_result: dict) -> KeyPoints:
if not inference_result.get("predictions"):
return cls.empty()
xy = []
confidence = []
class_id = []
class_names = []
for prediction in inference_result["predictions"]:
prediction_xy = []
prediction_confidence = []
for keypoint in prediction["keypoints"]:
prediction_xy.append([keypoint["x"], keypoint["y"]])
prediction_confidence.append(keypoint["confidence"])
xy.append(prediction_xy)
confidence.append(prediction_confidence)
class_id.append(prediction["class_id"])
class_names.append(prediction["class"])
data = {CLASS_NAME_DATA_FIELD: np.array(class_names)}
return cls(
xy=np.array(xy, dtype=np.float32),
confidence=np.array(confidence, dtype=np.float32),
class_id=np.array(class_id, dtype=int),
data=data,
)
@classmethod
def from_ultralytics(cls, ultralytics_results) -> KeyPoints:
"""