Repository: roboflow/supervision
File: supervision/keypoint/core.py (Python)

Code before:
        confidence = ultralytics_results.keypoints.conf.cpu().numpy()
        data = {CLASS_NAME_DATA_FIELD: class_names}
        return cls(xy, class_id, confidence, data)

    @classmethod
    def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints:
        """
        Create a KeyPoints instance from a YOLO NAS results.

        Args:
            yolo_nas_results (ImagePoseEstimationPrediction):
                The output object from YOLO NAS.

        Returns:
            KeyPoints: A new KeyPoints object.

        Example:
            ```python
            import cv2
            import supervision as sv
            import super_gradients

            image = cv2.imread(<SOURCE_IMAGE_PATH>)

            yolo_nas = super_gradients.training.models.get(
                "yolo_nas_pose_s", pretrained_weights="coco_pose").to("cuda")

            results = yolo_nas.predict(image, conf=0.1)
            keypoints = sv.KeyPoints.from_yolo_nas(results)
            ```
        """
        if len(yolo_nas_results.prediction.poses) == 0:
            return cls.empty()

        xy = yolo_nas_results.prediction.poses[:, :, :2]
        confidence = yolo_nas_results.prediction.poses[:, :, 2]
        class_id = [0] * len(xy)
        data = {CLASS_NAME_DATA_FIELD: ["person" for _ in xy]}

        return cls(
            xy=xy,
            confidence=confidence,
            class_id=np.array(class_id),
            data=data,
        )

    def __getitem__(
        self, index: Union[int, slice, List[int], np.ndarray, str]
    ) -> Union["KeyPoints", List, np.ndarray, None]:
        """
        Get a subset of the KeyPoints object or access an item from its data field.

Review comment (bug):
Using `to("cuda")` seems quite risky. It will crash if someone copies the snippet in an environment without a GPU. I would add `"cuda" if torch.cuda.is_available() else "cpu"` to the example. Just to be safe.



Code after:
        confidence = ultralytics_results.keypoints.conf.cpu().numpy()
        data = {CLASS_NAME_DATA_FIELD: class_names}
        return cls(xy, class_id, confidence, data)

    @classmethod
    def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints:
        """
        Create a KeyPoints instance from a YOLO NAS results.

        Args:
            yolo_nas_results (ImagePoseEstimationPrediction):
                The output object from YOLO NAS.

        Returns:
            KeyPoints: A new KeyPoints object.

        Example:
            ```python
            import cv2
            import torch
            import supervision as sv
            import super_gradients

            image = cv2.imread(<SOURCE_IMAGE_PATH>)

            device = "cuda" if torch.cuda.is_available() else "cpu"
            yolo_nas = super_gradients.training.models.get(
                "yolo_nas_pose_s", pretrained_weights="coco_pose").to(device)

            results = yolo_nas.predict(image, conf=0.1)
            keypoints = sv.KeyPoints.from_yolo_nas(results)
            ```
        """
        if len(yolo_nas_results.prediction.poses) == 0:
            return cls.empty()

        xy = yolo_nas_results.prediction.poses[:, :, :2]
        confidence = yolo_nas_results.prediction.poses[:, :, 2]

        # yolo_nas_results treats params differently.
        # prediction.labels may not exist, whereas class_names might be None
        if hasattr(yolo_nas_results.prediction, "labels"):
            class_id = yolo_nas_results.prediction.labels  # np.array[int]
        else:
            class_id = None

        data = {}
        if class_id is not None and yolo_nas_results.class_names is not None:
            class_names = []
            for c_id in class_id:
                name = yolo_nas_results.class_names[c_id]  # tuple[str]