diff --git a/supervision/__init__.py b/supervision/__init__.py index c1b9d2c6..60b48a3d 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -78,7 +78,7 @@ from supervision.detection.utils import ( xcycwh_to_xyxy, xywh_to_xyxy, xyxy_to_polygons, - xyxy_xywh, + xyxy_to_xywh, ) from supervision.draw.color import Color, ColorPalette from supervision.draw.utils import ( @@ -227,5 +227,5 @@ __all__ = [ "xcycwh_to_xyxy", "xywh_to_xyxy", "xyxy_to_polygons", - "xyxy_xywh", + "xyxy_to_xywh", ] diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index dfc66bab..61d0c0cb 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -321,7 +321,7 @@ def xywh_to_xyxy(xywh: np.ndarray) -> np.ndarray: return xyxy -def xyxy_xywh(xyxy: np.ndarray) -> np.ndarray: +def xyxy_to_xywh(xyxy: np.ndarray) -> np.ndarray: """ Converts bounding box coordinates from `(x_min, y_min, x_max, y_max)` format to `(x, y, width, height)` format. @@ -345,7 +345,7 @@ def xyxy_xywh(xyxy: np.ndarray) -> np.ndarray: [15, 25, 50, 70] ]) - sv.xyxy_xywh(xyxy=xyxy) + sv.xyxy_to_xywh(xyxy=xyxy) # array([ # [10, 20, 30, 40], # [15, 25, 35, 45] diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py index 4766ae73..dca0334e 100644 --- a/supervision/keypoint/core.py +++ b/supervision/keypoint/core.py @@ -526,53 +526,49 @@ class KeyPoints: Example: ```python - import requests - import torch from PIL import Image + import requests + import supervision as sv + import torch from transformers import ( AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation, ) - import supervision as sv - device = "cuda" if torch.cuda.is_available() else "cpu" image = Image.open() - person_image_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365") - person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device) + DETECTION_MODEL_ID = "PekingU/rtdetr_r50vd_coco_o365" - inputs = person_image_processor(images=image, return_tensors="pt").to(device) + detection_processor = AutoProcessor.from_pretrained(DETECTION_MODEL_ID, use_fast=True) + detection_model = RTDetrForObjectDetection.from_pretrained(DETECTION_MODEL_ID, device_map=DEVICE) + + inputs = detection_processor(images=frame, return_tensors="pt").to(DEVICE) with torch.no_grad(): - outputs = person_model(**inputs) + outputs = detection_model(**inputs) - results = person_image_processor.post_process_object_detection( - outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3 - ) - result = results[0] # take first image results - detections = sv.Detections.from_transformers(result) - person_detections_xywh = sv.xyxy_xywh(detections[detections.class_id == 0].xyxy) + target_size = torch.tensor([(frame.height, frame.width)]) + results = detection_processor.post_process_object_detection( + outputs, target_sizes=target_size, threshold=0.3) - image_processor = AutoProcessor.from_pretrained("usyd-community/vitpose-base-simple") - model = VitPoseForPoseEstimation.from_pretrained( - "usyd-community/vitpose-base-simple", device_map=device - ) + detections = sv.Detections.from_transformers(results[0]) + boxes = sv.xyxy_to_xywh(detections[detections.class_id == 0].xyxy) - inputs = image_processor(image, boxes=[person_detections_xywh], return_tensors="pt").to( - device - ) + POSE_ESTIMATION_MODEL_ID = "usyd-community/vitpose-base-simple" + + pose_estimation_processor = AutoProcessor.from_pretrained(POSE_ESTIMATION_MODEL_ID) + pose_estimation_model = VitPoseForPoseEstimation.from_pretrained( + POSE_ESTIMATION_MODEL_ID, device_map=DEVICE) + + inputs = pose_estimation_processor(frame, boxes=[boxes], return_tensors="pt").to(DEVICE) with torch.no_grad(): - outputs = model(**inputs) - - pose_results = image_processor.post_process_pose_estimation( - outputs, boxes=[person_detections_xywh] - )[0] - - keypoints = sv.KeyPoints.from_transformers(pose_results) + outputs = pose_estimation_model(**inputs) + results = pose_estimation_processor.post_process_pose_estimation(outputs, boxes=[boxes]) + key_point = sv.KeyPoints.from_transformers(results[0]) ``` """ # noqa: E501 // docs diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index d93c72c8..ed48ec86 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -21,6 +21,7 @@ from supervision.detection.utils import ( scale_boxes, xcycwh_to_xyxy, xywh_to_xyxy, + xyxy_to_xywh, ) TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) @@ -1381,6 +1382,29 @@ def test_xywh_to_xyxy(xywh: np.ndarray, expected_result: np.ndarray) -> None: np.testing.assert_array_equal(result, expected_result) +@pytest.mark.parametrize( + "xyxy, expected_result", + [ + (np.array([[10, 20, 40, 60]]), np.array([[10, 20, 30, 40]])), # standard case + (np.array([[0, 0, 0, 0]]), np.array([[0, 0, 0, 0]])), # zero size bounding box + ( + np.array([[50, 50, 150, 150]]), + np.array([[50, 50, 100, 100]]), + ), # large bounding box + ( + np.array([[-10, -20, 20, 20]]), + np.array([[-10, -20, 30, 40]]), + ), # negative coordinates + (np.array([[50, 50, 50, 80]]), np.array([[50, 50, 0, 30]])), # zero width + (np.array([[50, 50, 70, 50]]), np.array([[50, 50, 20, 0]])), # zero height + (np.array([]).reshape(0, 4), np.array([]).reshape(0, 4)), # empty array + ], +) +def test_xyxy_to_xywh(xyxy: np.ndarray, expected_result: np.ndarray) -> None: + result = xyxy_to_xywh(xyxy) + np.testing.assert_array_equal(result, expected_result) + + @pytest.mark.parametrize( "xcycwh, expected_result", [