diff --git a/supervision/key_points/core.py b/supervision/key_points/core.py index fb37e628..d8214170 100644 --- a/supervision/key_points/core.py +++ b/supervision/key_points/core.py @@ -88,12 +88,12 @@ class KeyPoints: key_points = sv.KeyPoints.from_mediapipe( pose_landmarker_result, (image_width, image_height)) ``` - + === "Transformers" - + Use [`sv.KeyPoints.from_transformers`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.from_transformers) method, which accepts [ViTPose](https://huggingface.co/docs/transformers/en/model_doc/vitpose) result. - + ```python from PIL import Image import requests @@ -650,7 +650,7 @@ class KeyPoints: return cls.empty() def __getitem__( - self, index: int | slice | list[int] | np.ndarray | tuple | str + self, index: int | slice | list[int] | np.ndarray | tuple | str ) -> KeyPoints | np.ndarray | list | None: if isinstance(index, str): return self.data.get(index) @@ -674,22 +674,18 @@ class KeyPoints: j = np.flatnonzero(j) if ( - isinstance(i, (list, np.ndarray)) - and isinstance(j, (list, np.ndarray)) - and not np.isscalar(i) - and not np.isscalar(j) + isinstance(i, (list, np.ndarray)) + and isinstance(j, (list, np.ndarray)) + and not np.isscalar(i) + and not np.isscalar(j) ): i, j = np.ix_(i, j) xy_selected = self.xy[i, j] - conf_selected = ( - self.confidence[i, j] if self.confidence is not None else None - ) + conf_selected = self.confidence[i, j] if self.confidence is not None else None - class_id_selected = ( - self.class_id[i] if self.class_id is not None else None - ) + class_id_selected = self.class_id[i] if self.class_id is not None else None data_selected = get_data_item(self.data, i) @@ -699,13 +695,13 @@ class KeyPoints: conf_selected = conf_selected.reshape(1, 1) elif xy_selected.ndim == 2: if np.isscalar(index[0]) or ( - isinstance(index[0], np.ndarray) and index[0].ndim == 0 + isinstance(index[0], np.ndarray) and index[0].ndim == 0 ): xy_selected = xy_selected[np.newaxis, ...] if conf_selected is not None: conf_selected = conf_selected[np.newaxis, ...] elif np.isscalar(index[1]) or ( - isinstance(index[1], np.ndarray) and index[1].ndim == 0 + isinstance(index[1], np.ndarray) and index[1].ndim == 0 ): xy_selected = xy_selected[:, np.newaxis, :] if conf_selected is not None: diff --git a/test/key_points/test_core.py b/test/key_points/test_core.py index f12cc2fd..5a8244cd 100644 --- a/test/key_points/test_core.py +++ b/test/key_points/test_core.py @@ -1,33 +1,21 @@ +from contextlib import nullcontext as DoesNotRaise + import numpy as np import pytest -from contextlib import nullcontext as DoesNotRaise + from supervision.key_points.core import KeyPoints from test.test_utils import mock_key_points KEY_POINTS = mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], - - [[20, 21], - [22, 23], - [24, 25], - [26, 27], - [28, 29]] + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], + [[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], [0.7, 0.9, 0.3, 0.4, 0.0], - [0.1, 0.6, 0.8, 0.2, 0.7] + [0.1, 0.6, 0.8, 0.2, 0.7], ], class_id=[0, 1, 2], ) @@ -52,16 +40,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, slice(0, 1), mock_key_points( - xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]] - ], - confidence=[ - [0.8, 0.2, 0.6, 0.1, 0.5] - ], + xy=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]], + confidence=[[0.8, 0.2, 0.6, 0.1, 0.5]], class_id=[0], ), DoesNotRaise(), @@ -71,17 +51,8 @@ KEY_POINTS = mock_key_points( slice(0, 2), mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], @@ -95,16 +66,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, 0, mock_key_points( - xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]] - ], - confidence=[ - [0.8, 0.2, 0.6, 0.1, 0.5] - ], + xy=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]], + confidence=[[0.8, 0.2, 0.6, 0.1, 0.5]], class_id=[0], ), DoesNotRaise(), @@ -113,16 +76,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, -1, mock_key_points( - xy=[ - [[20, 21], - [22, 23], - [24, 25], - [26, 27], - [28, 29]] - ], - confidence=[ - [0.1, 0.6, 0.8, 0.2, 0.7] - ], + xy=[[[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]]], + confidence=[[0.1, 0.6, 0.8, 0.2, 0.7]], class_id=[2], ), DoesNotRaise(), @@ -132,17 +87,8 @@ KEY_POINTS = mock_key_points( [0, 1], mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], @@ -157,17 +103,8 @@ KEY_POINTS = mock_key_points( np.array([0, 1]), mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], @@ -182,17 +119,8 @@ KEY_POINTS = mock_key_points( [True, True, False], mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], @@ -207,17 +135,8 @@ KEY_POINTS = mock_key_points( np.array([True, True, False]), mock_key_points( xy=[ - [[0, 1], - [2, 3], - [4, 5], - [6, 7], - [8, 9]], - - [[10, 11], - [12, 13], - [14, 15], - [16, 17], - [18, 19]], + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + [[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]], ], confidence=[ [0.8, 0.2, 0.6, 0.1, 0.5], @@ -237,18 +156,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), slice(0, 1)), mock_key_points( - xy=[ - [[0, 1]], - - [[10, 11]], - - [[20, 21]] - ], - confidence=[ - [0.8], - [0.7], - [0.1] - ], + xy=[[[0, 1]], [[10, 11]], [[20, 21]]], + confidence=[[0.8], [0.7], [0.1]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -257,21 +166,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), slice(0, 2)), mock_key_points( - xy=[ - [[0, 1], - [2, 3]], - - [[10, 11], - [12, 13]], - - [[20, 21], - [22, 23]] - ], - confidence=[ - [0.8, 0.2], - [0.7, 0.9], - [0.1, 0.6] - ], + xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]], + confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -280,18 +176,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), 0), mock_key_points( - xy=[ - [[0, 1]], - - [[10, 11]], - - [[20, 21]] - ], - confidence=[ - [0.8], - [0.7], - [0.1] - ], + xy=[[[0, 1]], [[10, 11]], [[20, 21]]], + confidence=[[0.8], [0.7], [0.1]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -300,18 +186,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), -1), mock_key_points( - xy=[ - [[8, 9]], - - [[18, 19]], - - [[28, 29]] - ], - confidence=[ - [0.5], - [0.0], - [0.7] - ], + xy=[[[8, 9]], [[18, 19]], [[28, 29]]], + confidence=[[0.5], [0.0], [0.7]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -320,21 +196,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), [0, 1]), mock_key_points( - xy=[ - [[0, 1], - [2, 3]], - - [[10, 11], - [12, 13]], - - [[20, 21], - [22, 23]] - ], - confidence=[ - [0.8, 0.2], - [0.7, 0.9], - [0.1, 0.6] - ], + xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]], + confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -343,21 +206,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), np.array([0, 1])), mock_key_points( - xy=[ - [[0, 1], - [2, 3]], - - [[10, 11], - [12, 13]], - - [[20, 21], - [22, 23]] - ], - confidence=[ - [0.8, 0.2], - [0.7, 0.9], - [0.1, 0.6] - ], + xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]], + confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -366,21 +216,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), [True, True, False, False, False]), mock_key_points( - xy=[ - [[0, 1], - [2, 3]], - - [[10, 11], - [12, 13]], - - [[20, 21], - [22, 23]] - ], - confidence=[ - [0.8, 0.2], - [0.7, 0.9], - [0.1, 0.6] - ], + xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]], + confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -389,21 +226,8 @@ KEY_POINTS = mock_key_points( KEY_POINTS, (slice(None), np.array([True, True, False, False, False])), mock_key_points( - xy=[ - [[0, 1], - [2, 3]], - - [[10, 11], - [12, 13]], - - [[20, 21], - [22, 23]] - ], - confidence=[ - [0.8, 0.2], - [0.7, 0.9], - [0.1, 0.6] - ], + xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]], + confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]], class_id=[0, 1, 2], ), DoesNotRaise(), @@ -422,7 +246,7 @@ KEY_POINTS = mock_key_points( ), DoesNotRaise(), ), # select the first anchor from the first skeleton by index -( + ( KEY_POINTS, (0, -1), mock_key_points(