diff --git a/supervision/metrics/detection.py b/supervision/metrics/detection.py index d0e85e63..e7234973 100644 --- a/supervision/metrics/detection.py +++ b/supervision/metrics/detection.py @@ -84,8 +84,12 @@ class ConfusionMatrix: prediction_tensors = [] target_tensors = [] for prediction, target in zip(predictions, targets): - prediction_tensors.append(cls.convert_detections_to_tensor(prediction)) - target_tensors.append(cls.convert_detections_to_tensor(target)) + prediction_tensors.append( + cls.detections_to_tensor(prediction, with_confidence=True) + ) + target_tensors.append( + cls.detections_to_tensor(target, with_confidence=False) + ) return cls.from_tensors( predictions=prediction_tensors, targets=target_tensors, @@ -95,15 +99,24 @@ class ConfusionMatrix: ) @classmethod - def convert_detections_to_tensor(cls, detections: Detections) -> np.ndarray: + def detections_to_tensor( + cls, detections: Detections, with_confidence: bool = False + ) -> np.ndarray: + if detections.class_id is None: + raise ValueError( + "ConfusionMatrix can only be calculated for Detections with class_id" + ) + arrays_to_concat = [detections.xyxy, np.expand_dims(detections.class_id, 1)] - if detections.confidence is not None: + + if with_confidence: + if detections.confidence is None: + raise ValueError( + "ConfusionMatrix can only be calculated for Detections with confidence" + ) arrays_to_concat.append(np.expand_dims(detections.confidence, 1)) - return np.concatenate( - arrays_to_concat, - axis=1, - ) + return np.concatenate(arrays_to_concat, axis=1) @classmethod def from_tensors( diff --git a/test/metrics/test_detection.py b/test/metrics/test_detection.py index 4bd72cf4..0ecab3f0 100644 --- a/test/metrics/test_detection.py +++ b/test/metrics/test_detection.py @@ -6,6 +6,7 @@ import pytest from supervision.detection.core import Detections from supervision.metrics.detection import ConfusionMatrix +from test.utils import mock_detections CLASSES = np.arange(80) NUM_CLASSES = len(CLASSES) @@ -119,26 +120,58 @@ BAD_CONF_MATRIX = worsen_ideal_conf_matrix( @pytest.mark.parametrize( - "detections, exception", + "detections, with_confidence, expected_result, exception", [ ( - DETECTIONS, + Detections.empty(), + False, + np.empty((0, 5), dtype=np.float32), DoesNotRaise(), - ) + ), # empty detections; no confidence + ( + Detections.empty(), + True, + np.empty((0, 6), dtype=np.float32), + DoesNotRaise(), + ), # empty detections; with confidence + ( + mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]), + False, + np.array([[0, 0, 10, 10, 0]], dtype=np.float32), + DoesNotRaise(), + ), # single detection; no confidence + ( + mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]), + True, + np.array([[0, 0, 10, 10, 0, 0.5]], dtype=np.float32), + DoesNotRaise(), + ), # single detection; with confidence + ( + mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]], class_id=[0, 1], confidence=[0.5, 0.2]), + False, + np.array([[0, 0, 10, 10, 0], [0, 0, 20, 20, 1]], dtype=np.float32), + DoesNotRaise(), + ), # multiple detections; no confidence + ( + mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]], class_id=[0, 1], confidence=[0.5, 0.2]), + True, + np.array([[0, 0, 10, 10, 0, 0.5], [0, 0, 20, 20, 1, 0.2]], dtype=np.float32), + DoesNotRaise(), + ), # multiple detections; with confidence ], ) -def test_convert_detections_to_tensor( - detections, - exception: Exception, +def test_detections_to_tensor( + detections: Detections, + with_confidence: bool, + expected_result: Optional[np.ndarray], + exception: Exception ): with exception: - result = ConfusionMatrix.convert_detections_to_tensor( + result = ConfusionMatrix.detections_to_tensor( detections=detections, + with_confidence=with_confidence ) - - assert np.array_equal(result[:, :4], detections.xyxy) - assert np.array_equal(result[:, 4], detections.class_id) - assert np.array_equal(result[:, 5], detections.confidence) + assert np.array_equal(result, expected_result) @pytest.mark.parametrize(