From 6dde70eb0c7ffcfa97191b9276bbbc77c82471ce Mon Sep 17 00:00:00 2001 From: kirilllzaitsev Date: Wed, 26 Jul 2023 10:46:31 +0200 Subject: [PATCH] move fixing of class_ids to load_pascal_voc_annotations --- supervision/dataset/core.py | 4 ---- supervision/dataset/formats/pascal_voc.py | 4 ++++ 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index d40728e4..15d92aa0 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -241,10 +241,6 @@ class DetectionDataset(BaseDataset): force_masks=force_masks, ) - for annotation in annotations.values(): - class_id = [classes.index(class_name) for class_name in annotation.class_id] - annotation.class_id = np.array(class_id) - return DetectionDataset(classes=classes, images=images, annotations=annotations) @classmethod diff --git a/supervision/dataset/formats/pascal_voc.py b/supervision/dataset/formats/pascal_voc.py index 90501be3..8f3fb090 100644 --- a/supervision/dataset/formats/pascal_voc.py +++ b/supervision/dataset/formats/pascal_voc.py @@ -203,6 +203,10 @@ def load_pascal_voc_annotations( classes = list(set(classes)) + for annotation in annotations.values(): + class_id = [classes.index(class_name) for class_name in annotation.class_id] + annotation.class_id = np.array(class_id) + return classes, images, annotations