Merge pull request #72 from roboflow/feature/add_pascal_voc_xml_load_…

…support

feature/add_pascal_voc_xml_load_support

docs/changelog.md

@@ -1,3 +1,11 @@
+### 0.6.0 <small>April 19, 2023</small>
+
+- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format. 
+- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
+- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
+- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
+- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
+
 ### 0.5.2 <small>April 13, 2023</small>
 
 - Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.

supervision/dataset/core.py

@@ -1,12 +1,18 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import cv2
 import numpy as np
 
-from supervision.dataset.formats.pascal_voc import detections_to_pascal_voc
+from supervision.dataset.formats.pascal_voc import (
+    detections_to_pascal_voc,
+    load_pascal_voc_annotations,
+)
 from supervision.detection.core import Detections
+from supervision.file import list_files_with_extensions
 
 
 @dataclass
@@ -76,3 +82,47 @@ def as_pascal_voc(
 
                 with open(annotations_path / f"{annotation_name}.xml", "w") as f:
                     f.write(pascal_voc_xml)
+
+    @classmethod
+    def from_pascal_voc(
+        cls, images_directory_path: str, annotations_directory_path: str
+    ) -> Dataset:
+        """
+        Creates a Dataset instance from PASCAL VOC formatted data.
+
+        Args:
+            images_directory_path (str): The path to the directory containing the images.
+            annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
+
+        Returns:
+            Dataset: A Dataset instance containing the loaded images and annotations.
+        """
+        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"]
+        )
+
+        raw_annotations: List[Tuple[str, Detections, List[str]]] = [
+            load_pascal_voc_annotations(annotation_path=str(annotation_path))
+            for annotation_path in annotation_paths
+        ]
+
+        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 Dataset(classes=classes, images=images, annotations=annotations)

supervision/dataset/formats/pascal_voc.py

@@ -1,6 +1,6 @@
 from typing import List, Optional, Tuple
 from xml.dom.minidom import parseString
-from xml.etree.ElementTree import Element, SubElement, tostring
+from xml.etree.ElementTree import Element, SubElement, parse, tostring
 
 import numpy as np
 
@@ -132,3 +132,40 @@ def detections_to_pascal_voc(
     xml_string = parseString(tostring(annotation)).toprettyxml(indent="  ")
 
     return xml_string
+
+
+def load_pascal_voc_annotations(
+    annotation_path: str,
+) -> Tuple[str, Detections, List[str]]:
+    """
+    Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names.
+
+    Args:
+        annotation_path (str): The path to the PASCAL VOC XML annotations file.
+
+    Returns:
+        Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections.
+    """
+    tree = parse(annotation_path)
+    root = tree.getroot()
+
+    image_name = root.find("filename").text
+
+    xyxy = []
+    class_names = []
+    for obj in root.findall("object"):
+        class_name = obj.find("name").text
+        class_names.append(class_name)
+
+        bbox = obj.find("bndbox")
+        x1 = int(bbox.find("xmin").text)
+        y1 = int(bbox.find("ymin").text)
+        x2 = int(bbox.find("xmax").text)
+        y2 = int(bbox.find("ymax").text)
+
+        xyxy.append([x1, y1, x2, y2])
+
+    xyxy = np.array(xyxy)
+    detections = Detections(xyxy=xyxy)
+
+    return image_name, detections, class_names