From c9598ef732c62010530a5dcad2e4a91056614202 Mon Sep 17 00:00:00 2001 From: arthurcerveira Date: Fri, 26 Jul 2024 10:46:14 -0300 Subject: [PATCH] refactor: use image lazy loading for benchmarks --- supervision/metrics/detection.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/supervision/metrics/detection.py b/supervision/metrics/detection.py index 99a11322..bbbfa711 100644 --- a/supervision/metrics/detection.py +++ b/supervision/metrics/detection.py @@ -391,11 +391,10 @@ class ConfusionMatrix: ``` """ predictions, targets = [], [] - for img_name, img in dataset.images.items(): - predictions_batch = callback(img) + for _, image, annotation in dataset: + predictions_batch = callback(image) predictions.append(predictions_batch) - targets_batch = dataset.annotations[img_name] - targets.append(targets_batch) + targets.append(annotation) return cls.from_detections( predictions=predictions, targets=targets, @@ -604,11 +603,10 @@ class MeanAveragePrecision: ``` """ predictions, targets = [], [] - for img_name, img in dataset.images.items(): - predictions_batch = callback(img) + for _, image, annotation in dataset: + predictions_batch = callback(image) predictions.append(predictions_batch) - targets_batch = dataset.annotations[img_name] - targets.append(targets_batch) + targets.append(annotation) return cls.from_detections( predictions=predictions, targets=targets,