roboflow--supervision/tests/dataset/formats/test_pascal_voc.py

from contextlib import ExitStack as DoesNotRaise
from pathlib import Path

import cv2
import numpy as np
import pytest
from defusedxml import ElementTree

from supervision.dataset.core import DetectionDataset
from supervision.dataset.formats.pascal_voc import (
    detections_from_xml_obj,
    detections_to_pascal_voc,
    load_pascal_voc_annotations,
    object_to_pascal_voc,
    parse_polygon_points,
    save_pascal_voc_annotations,
)
from tests.helpers import _create_detections


def are_xml_elements_equal(elem1, elem2) -> bool:
    if (
        elem1.tag != elem2.tag
        or elem1.attrib != elem2.attrib
        or elem1.text != elem2.text
        or len(elem1) != len(elem2)
    ):
        return False

    for child1, child2 in zip(elem1, elem2):
        if not are_xml_elements_equal(child1, child2):
            return False

    return True


@pytest.mark.parametrize(
    ("xyxy", "name", "polygon", "expected_result", "exception"),
    [
        pytest.param(
            np.array([0, 0, 10, 10]),
            "test",
            None,
            ElementTree.fromstring(
                """<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
                <xmax>11</xmax><ymax>11</ymax></bndbox></object>"""
            ),
            DoesNotRaise(),
            id="bbox_only",
        ),
        pytest.param(
            np.array([0, 0, 10, 10]),
            "test",
            np.array([[0, 0], [10, 0], [10, 10], [0, 10]]),
            ElementTree.fromstring(
                """<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
                <xmax>11</xmax><ymax>11</ymax>
                </bndbox><polygon><x1>1</x1><y1>1</y1><x2>11</x2>
                <y2>1</y2><x3>11</x3><y3>11</y3><x4>1</x4><y4>11</y4>
                </polygon></object>"""
            ),
            DoesNotRaise(),
            id="bbox_and_polygon",
        ),
    ],
)
def test_object_to_pascal_voc(
    xyxy: np.ndarray,
    name: str,
    polygon: np.ndarray | None,
    expected_result,
    exception: Exception,
) -> None:
    with exception:
        result = object_to_pascal_voc(xyxy=xyxy, name=name, polygon=polygon)
        assert are_xml_elements_equal(result, expected_result)


def test_object_to_pascal_voc_does_not_mutate_inputs():
    """Serializing an object must not write the 1-index offset back into the inputs."""
    xyxy = np.array([10, 20, 30, 40], dtype=np.float32)
    polygon = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)

    object_to_pascal_voc(xyxy=xyxy, name="test", polygon=polygon)

    assert np.array_equal(xyxy, np.array([10, 20, 30, 40], dtype=np.float32))
    assert np.array_equal(
        polygon, np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32)
    )


def test_object_to_pascal_voc_does_not_mutate_view_input():
    """Mutation guard holds when xyxy is a NumPy row-view (the actual bug scenario)."""
    base = np.array([[10, 20, 30, 40]], dtype=np.float32)
    xyxy_view = base[0]  # row-view, shares memory with base

    object_to_pascal_voc(xyxy=xyxy_view, name="test", polygon=None)

    assert np.array_equal(base[0], np.array([10, 20, 30, 40], dtype=np.float32)), (
        "object_to_pascal_voc mutated the source array via a view"
    )


def test_detections_to_pascal_voc_does_not_mutate_detections():
    """Exporting detections must not shift the source xyxy, and must be repeatable."""
    detections = _create_detections(xyxy=[[10, 20, 30, 40]], class_id=[0])
    expected_xyxy = detections.xyxy.copy()

    first = detections_to_pascal_voc(
        detections, classes=["test"], filename="image.jpg", image_shape=(100, 100, 3)
    )
    second = detections_to_pascal_voc(
        detections, classes=["test"], filename="image.jpg", image_shape=(100, 100, 3)
    )

    assert np.array_equal(detections.xyxy, expected_xyxy)
    assert first == second


@pytest.mark.parametrize(
    ("polygon_element", "expected_result", "exception"),
    [
        pytest.param(
            ElementTree.fromstring(
                """<polygon><x1>0</x1><y1>0</y1><x2>10</x2><y2>0</y2><x3>10</x3>
                    <y3>10</y3><x4>0</x4><y4>10</y4></polygon>"""
            ),
            np.array([[0, 0], [10, 0], [10, 10], [0, 10]]),
            DoesNotRaise(),
            id="standard_polygon",
        )
    ],
)
def test_parse_polygon_points(
    polygon_element,
    expected_result: list[list],
    exception,
) -> None:
    with exception:
        result = parse_polygon_points(polygon_element)
        assert np.array_equal(result, expected_result)


ONE_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
</bndbox></object><object><name>test</name><bndbox><xmin>11</xmin><ymin>11</ymin>
<xmax>21</xmax><ymax>21</ymax></bndbox></object></annotation>"""


ONE_CLASS_ONE_BBOX = """<annotation><object><name>test</name><bndbox>
<xmin>1</xmin><ymin>1</ymin><xmax>11</xmax><ymax>11</ymax></bndbox></object>
</annotation>"""


N_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
</bndbox></object><object><name>test</name><bndbox>
<xmin>21</xmin><ymin>31</ymin><xmax>31</xmax><ymax>41</ymax></bndbox>
</object><object><name>test2</name><bndbox><xmin>
11</xmin><ymin>11</ymin><xmax>21</xmax><ymax>
21</ymax></bndbox></object></annotation>"""

NO_DETECTIONS = """<annotation></annotation>"""
MIXED_POLYGON_AND_BOX = """<annotation><object><name>test</name><bndbox>
<xmin>1</xmin><ymin>1</ymin><xmax>11</xmax><ymax>11</ymax></bndbox>
<polygon><x1>1</x1><y1>1</y1><x2>11</x2><y2>1</y2><x3>11</x3><y3>11</y3>
<x4>1</x4><y4>11</y4></polygon></object><object><name>test</name><bndbox>
<xmin>11</xmin><ymin>11</ymin><xmax>21</xmax><ymax>21</ymax></bndbox></object>
</annotation>"""


@pytest.mark.parametrize(
    (
        "xml_string",
        "classes",
        "resolution_wh",
        "force_masks",
        "expected_result",
        "exception",
    ),
    [
        pytest.param(
            ONE_CLASS_ONE_BBOX,
            ["test"],
            (100, 100),
            False,
            _create_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
            DoesNotRaise(),
            id="one_class_one_bbox",
        ),
        pytest.param(
            ONE_CLASS_N_BBOX,
            ["test"],
            (100, 100),
            False,
            _create_detections(
                xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), class_id=[0, 0]
            ),
            DoesNotRaise(),
            id="one_class_n_bbox",
        ),
        pytest.param(
            N_CLASS_N_BBOX,
            ["test", "test2"],
            (100, 100),
            False,
            _create_detections(
                xyxy=np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
                class_id=[0, 0, 1],
            ),
            DoesNotRaise(),
            id="n_class_n_bbox",
        ),
        pytest.param(
            NO_DETECTIONS,
            [],
            (100, 100),
            False,
            _create_detections(xyxy=np.empty((0, 4)), class_id=[]),
            DoesNotRaise(),
            id="no_detections",
        ),
    ],
)
def test_detections_from_xml_obj(
    xml_string, classes, resolution_wh, force_masks, expected_result, exception
) -> None:
    with exception:
        root = ElementTree.fromstring(xml_string)
        result, _ = detections_from_xml_obj(root, classes, resolution_wh, force_masks)
        assert result == expected_result


@pytest.mark.parametrize("force_masks", [False, True])
def test_detections_from_xml_obj_mixed_polygon_and_bbox_masks_aligned(
    force_masks: bool,
) -> None:
    root = ElementTree.fromstring(MIXED_POLYGON_AND_BOX)
    detections, _ = detections_from_xml_obj(
        root=root,
        classes=["test"],
        resolution_wh=(30, 30),
        force_masks=force_masks,
    )

    assert detections.mask is not None
    assert detections.mask.shape == (2, 30, 30)
    assert detections.mask[0].any()
    assert not detections.mask[1].any()


def _write_voc_sample(
    images_dir: Path, annotations_dir: Path, stem: str, class_names: list[str]
) -> None:
    """Write one VOC image plus its bbox-only XML annotation to disk."""
    cv2.imwrite(str(images_dir / f"{stem}.png"), np.zeros((20, 20, 3), dtype=np.uint8))
    objects = "".join(
        f"<object><name>{name}</name><bndbox><xmin>1</xmin><ymin>1</ymin>"
        f"<xmax>10</xmax><ymax>10</ymax></bndbox></object>"
        for name in class_names
    )
    (annotations_dir / f"{stem}.xml").write_text(f"<annotation>{objects}</annotation>")


class TestLoadPascalVocDeterministicClasses:
    """Regression tests for deterministic VOC class ordering (DAT-03)."""

    def test_classes_sorted_within_file(self, tmp_path: Path) -> None:
        """Class names from one file are assigned ids in sorted, stable order."""
        images_dir = tmp_path / "images"
        images_dir.mkdir()
        annotations_dir = tmp_path / "annotations"
        annotations_dir.mkdir()
        _write_voc_sample(images_dir, annotations_dir, "img", ["zebra", "ant", "mango"])

        classes, _, _ = load_pascal_voc_annotations(
            images_directory_path=str(images_dir),
            annotations_directory_path=str(annotations_dir),
        )

        assert classes == sorted(classes) == ["ant", "mango", "zebra"]

    def test_repeated_loads_give_identical_class_ids(self, tmp_path: Path) -> None:
        """Two loads of the same multi-file VOC set produce identical class ids."""
        images_dir = tmp_path / "images"
        images_dir.mkdir()
        annotations_dir = tmp_path / "annotations"
        annotations_dir.mkdir()
        _write_voc_sample(images_dir, annotations_dir, "a_img", ["zebra"])
        _write_voc_sample(images_dir, annotations_dir, "b_img", ["ant", "mango"])

        first = load_pascal_voc_annotations(
            images_directory_path=str(images_dir),
            annotations_directory_path=str(annotations_dir),
        )
        second = load_pascal_voc_annotations(
            images_directory_path=str(images_dir),
            annotations_directory_path=str(annotations_dir),
        )

        assert first[0] == second[0]
        assert {p: d.class_id.tolist() for p, d in first[2].items()} == {
            p: d.class_id.tolist() for p, d in second[2].items()
        }


class TestSavePascalVocAnnotations:
    """save_pascal_voc_annotations: filesystem output contract."""

    def test_empty_dataset_creates_directory_and_no_xml_files(
        self, tmp_path: Path
    ) -> None:
        """Empty dataset produces no XML files; output directory is created."""
        dataset = DetectionDataset(classes=[], images=[], annotations={})
        out_dir = tmp_path / "annotations"

        save_pascal_voc_annotations(dataset, str(out_dir))

        assert out_dir.is_dir()
        assert list(out_dir.glob("*.xml")) == []

    def test_zero_detection_image_writes_xml_without_object_elements(
        self, tmp_path: Path
    ) -> None:
        """Image with no detections produces one XML file with no object elements."""
        from supervision.detection.core import Detections

        img_path = tmp_path / "img.jpg"
        cv2.imwrite(str(img_path), np.zeros((50, 50, 3), dtype=np.uint8))

        dataset = DetectionDataset(
            classes=["cat"],
            images=[str(img_path)],
            annotations={str(img_path): Detections.empty()},
        )
        out_dir = tmp_path / "annotations"

        save_pascal_voc_annotations(dataset, str(out_dir))

        xml_files = list(out_dir.glob("*.xml"))
        assert len(xml_files) == 1
        tree = ElementTree.parse(str(xml_files[0]))
        assert tree.findall("object") == []

    def test_show_progress_true_is_accepted_without_error(self, tmp_path: Path) -> None:
        """show_progress=True is accepted by the function without raising."""
        from supervision.detection.core import Detections

        img_path = tmp_path / "img.jpg"
        cv2.imwrite(str(img_path), np.zeros((50, 50, 3), dtype=np.uint8))

        dataset = DetectionDataset(
            classes=["cat"],
            images=[str(img_path)],
            annotations={str(img_path): Detections.empty()},
        )
        out_dir = tmp_path / "annotations"

        save_pascal_voc_annotations(dataset, str(out_dir), show_progress=True)

        assert out_dir.is_dir()