From 3723eae50c385fc2cbbf4e3cd97b68d5056e28ef Mon Sep 17 00:00:00 2001
From: kirilllzaitsev <kzaitsev88@gmail.com>
Date: Wed, 26 Jul 2023 10:21:17 +0200
Subject: [PATCH] update from_pascal_voc to match v2 loader

---
 supervision/dataset/core.py | 38 ++++++++++++-------------------------
 1 file changed, 12 insertions(+), 26 deletions(-)

diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py
index ce7e0edb..4bd80c1b 100644
--- a/supervision/dataset/core.py
+++ b/supervision/dataset/core.py
@@ -197,7 +197,10 @@ class DetectionDataset(BaseDataset):
 
     @classmethod
     def from_pascal_voc(
-        cls, images_directory_path: str, annotations_directory_path: str
+        cls,
+        images_directory_path: str,
+        annotations_directory_path: str,
+        force_masks: bool = False,
     ) -> DetectionDataset:
         """
         Creates a Dataset instance from PASCAL VOC formatted data.
@@ -231,34 +234,17 @@ class DetectionDataset(BaseDataset):
             ['dog', 'person']
             ```
         """
-        image_paths = list_files_with_extensions(
-            directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
-        )
-        annotation_paths = list_files_with_extensions(
-            directory=annotations_directory_path, extensions=["xml"]
+
+        classes, images, annotations = load_pascal_voc_annotations(
+            images_directory_path=images_directory_path,
+            annotations_directory_path=annotations_directory_path,
+            force_masks=force_masks,
         )
 
-        raw_annotations: List[Tuple[str, Detections, List[str]]] = [
-            load_pascal_voc_annotations(annotation_path=str(annotation_path))
-            for annotation_path in annotation_paths
-        ]
+        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)
 
-        classes = []
-        for annotation in raw_annotations:
-            classes.extend(annotation[2])
-        classes = list(set(classes))
-
-        for annotation in raw_annotations:
-            class_id = [classes.index(class_name) for class_name in annotation[2]]
-            annotation[1].class_id = np.array(class_id)
-
-        images = {
-            image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
-        }
-
-        annotations = {
-            image_name: detections for image_name, detections, _ in raw_annotations
-        }
         return DetectionDataset(classes=classes, images=images, annotations=annotations)
 
     @classmethod