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