From c8ea42af6df5084ddccbbbdd9cc18430b90a82a4 Mon Sep 17 00:00:00 2001 From: macc-n Date: Thu, 4 Apr 2024 17:32:41 +0200 Subject: [PATCH] Fixed obb key check in ultralytics_results --- supervision/detection/core.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 8c29cfc4..e378e4c8 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -208,22 +208,25 @@ class Detections: ``` """ # noqa: E501 // docs - if ultralytics_results.obb is not None: - class_id = ultralytics_results.obb.cls.cpu().numpy().astype(int) - class_names = np.array([ultralytics_results.names[i] for i in class_id]) - oriented_box_coordinates = ultralytics_results.obb.xyxyxyxy.cpu().numpy() - return cls( - xyxy=ultralytics_results.obb.xyxy.cpu().numpy(), - confidence=ultralytics_results.obb.conf.cpu().numpy(), - class_id=class_id, - tracker_id=ultralytics_results.obb.id.int().cpu().numpy() - if ultralytics_results.obb.id is not None - else None, - data={ - ORIENTED_BOX_COORDINATES: oriented_box_coordinates, - CLASS_NAME_DATA_FIELD: class_names, - }, - ) + if "obb" in ultralytics_results: + if ultralytics_results.obb is not None: + class_id = ultralytics_results.obb.cls.cpu().numpy().astype(int) + class_names = np.array([ultralytics_results.names[i] for i in class_id]) + oriented_box_coordinates = ( + ultralytics_results.obb.xyxyxyxy.cpu().numpy() + ) + return cls( + xyxy=ultralytics_results.obb.xyxy.cpu().numpy(), + confidence=ultralytics_results.obb.conf.cpu().numpy(), + class_id=class_id, + tracker_id=ultralytics_results.obb.id.int().cpu().numpy() + if ultralytics_results.obb.id is not None + else None, + data={ + ORIENTED_BOX_COORDINATES: oriented_box_coordinates, + CLASS_NAME_DATA_FIELD: class_names, + }, + ) class_id = ultralytics_results.boxes.cls.cpu().numpy().astype(int) class_names = np.array([ultralytics_results.names[i] for i in class_id])