Merge pull request #158 from hardikdava/dataset_merge

Dataset merge

supervision/dataset/core.py

@@ -23,7 +23,13 @@ from supervision.dataset.formats.yolo import (
     save_data_yaml,
     save_yolo_annotations,
 )
-from supervision.dataset.utils import save_dataset_images, train_test_split
+from supervision.dataset.utils import (
+    build_class_index_mapping,
+    map_detections_class_id,
+    merge_class_lists,
+    save_dataset_images,
+    train_test_split,
+)
 from supervision.detection.core import Detections
 from supervision.utils.file import list_files_with_extensions
 
@@ -76,6 +82,21 @@ class DetectionDataset(BaseDataset):
         for image_name, image in self.images.items():
             yield image_name, image, self.annotations.get(image_name, None)
 
+    def __eq__(self, other):
+        if not isinstance(other, DetectionDataset):
+            return False
+
+        if set(self.classes) != set(other.classes):
+            return False
+
+        for key in self.images:
+            if not np.array_equal(self.images[key], other.images[key]):
+                return False
+            if not self.annotations[key] == other.annotations[key]:
+                return False
+
+        return True
+
     def split(
         self, split_ratio=0.8, random_state=None, shuffle: bool = True
     ) -> Tuple[DetectionDataset, DetectionDataset]:
@@ -432,6 +453,54 @@ class DetectionDataset(BaseDataset):
                 info=info,
             )
 
+    @classmethod
+    def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset:
+        """
+        Merge a list of DetectionDataset objects into a single DetectionDataset object.
+
+        This method takes a list of DetectionDataset objects and combines their respective fields (`classes`, `images`,
+        `annotations`) into a single DetectionDataset object.
+
+        Args:
+            dataset_list (List[DetectionDataset]): A list of DetectionDataset objects to merge.
+
+        Returns:
+            (DetectionDataset): A single DetectionDataset object containing the merged data from the input list.
+
+        Example:
+            ```python
+            >>> from supervision import DetectionDataset
+
+            >>> ds_1 = DetectionDataset(...)
+            >>> ds_2 = DetectionDataset(...)
+
+            >>> merged_detection_dataset = DetectionDataset.merge([ds_1, ds_2])
+            ```
+        """
+        merged_images, merged_annotations = {}, {}
+        class_lists = [dataset.classes for dataset in dataset_list]
+        merged_classes = merge_class_lists(class_lists=class_lists)
+
+        for dataset in dataset_list:
+            class_index_mapping = build_class_index_mapping(
+                source_classes=dataset.classes, target_classes=merged_classes
+            )
+            for image_name, image, detections in dataset:
+                if image_name in merged_annotations:
+                    raise ValueError(
+                        f"Image name {image_name} is not unique across datasets."
+                    )
+
+                merged_images[image_name] = image
+                merged_annotations[image_name] = map_detections_class_id(
+                    source_to_target_mapping=class_index_mapping,
+                    detections=detections,
+                )
+
+        return cls(
+            classes=merged_classes, images=merged_images, annotations=merged_annotations
+        )
+
 
 @dataclass
 class ClassificationDataset(BaseDataset):

supervision/dataset/formats/coco.py

@@ -16,7 +16,6 @@ def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
     return [
         category["name"]
         for category in sorted(coco_categories, key=lambda category: category["id"])
-        if category["supercategory"] != "none"
     ]
 
 

supervision/dataset/utils.py

@@ -1,3 +1,4 @@
+import copy
 import os
 import random
 from pathlib import Path
@@ -6,6 +7,7 @@ from typing import Dict, List, Optional, Tuple, TypeVar
 import cv2
 import numpy as np
 
+from supervision.detection.core import Detections
 from supervision.detection.utils import (
     approximate_polygon,
     filter_polygons_by_area,
@@ -43,6 +45,51 @@ def approximate_mask_with_polygons(
     ]
 
 
+def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
+    unique_classes = set()
+
+    for class_list in class_lists:
+        for class_name in class_list:
+            unique_classes.add(class_name.lower())
+
+    return sorted(list(unique_classes))
+
+
+def build_class_index_mapping(
+    source_classes: List[str], target_classes: List[str]
+) -> Dict[int, int]:
+    index_mapping = {}
+
+    for i, class_name in enumerate(source_classes):
+        if class_name not in target_classes:
+            raise ValueError(
+                f"Class {class_name} not found in target classes. "
+                f"source_classes must be a subset of target_classes."
+            )
+        corresponding_index = target_classes.index(class_name)
+        index_mapping[i] = corresponding_index
+
+    return index_mapping
+
+
+def map_detections_class_id(
+    source_to_target_mapping: Dict[int, int], detections: Detections
+) -> Detections:
+    if detections.class_id is None:
+        raise ValueError("Detections must have class_id attribute.")
+    if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
+        raise ValueError(
+            "Detections class_id must be a subset of source_to_target_mapping keys."
+        )
+
+    detections_copy = copy.deepcopy(detections)
+    detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
+        detections_copy.class_id
+    )
+
+    return detections_copy
+
+
 def save_dataset_images(
     images_directory_path: str, images: Dict[str, np.ndarray]
 ) -> None:

test/dataset/formats/test_coco.py

@@ -10,7 +10,7 @@ from supervision.dataset.formats.coco import classes_to_coco_categories, coco_ca
 import numpy as np
 
 
-def generate_cock_coco_annotation(
+def mock_cock_coco_annotation(
     annotation_id: int = 0,
     image_id: int = 0,
     category_id: int = 0,
@@ -43,7 +43,7 @@ def generate_cock_coco_annotation(
                     "supercategory": "none"
                 }
             ],
-            [],
+            ["fashion-assistant"],
             DoesNotRaise()
         ),  # single coco category with supercategory == "none"
         (
@@ -60,7 +60,7 @@ def generate_cock_coco_annotation(
                 }
             ],
             [
-                "baseball cap"
+                "fashion-assistant", "baseball cap"
             ],
             DoesNotRaise()
         ),  # two coco categories; one with supercategory == "none" and one with supercategory != "none"
@@ -83,6 +83,7 @@ def generate_cock_coco_annotation(
                 }
             ],
             [
+                "fashion-assistant",
                 "baseball cap",
                 "hoodie"
             ],
@@ -107,6 +108,7 @@ def generate_cock_coco_annotation(
                 }
             ],
             [
+                "fashion-assistant",
                 "baseball cap",
                 "hoodie"
             ],
@@ -163,55 +165,55 @@ def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exce
         ),  # empty coco annotations
         (
             [
-                generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+                mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
             ],
             {
                 0: [
-                    generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
                 ]
             },
             DoesNotRaise()
         ),  # single coco annotation
         (
             [
-                generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
-                generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
+                mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
+                mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
             ],
             {
                 0: [
-                    generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
                 ],
                 1: [
-                    generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
+                    mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
                 ]
             },
             DoesNotRaise()
         ),  # two coco annotations
         (
             [
-                generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
-                generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
-                generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
-                generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
-                generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
-                generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
-                generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
+                mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
+                mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
+                mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
+                mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
+                mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
+                mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
+                mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
             ],
             {
                 0: [
-                    generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
                 ],
                 1: [
-                    generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
-                    generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
+                    mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
+                    mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
                 ],
                 2: [
-                    generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
+                    mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
                 ],
                 3: [
-                    generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
-                    generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
-                    generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
+                    mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
+                    mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
+                    mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
                 ]
             },
             DoesNotRaise()
@@ -240,7 +242,7 @@ def test_group_coco_annotations_by_image_id(
         ),  # empty image annotations
         (
             [
-                generate_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)
+                mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)
             ],
             (1000, 1000),
             False,
@@ -256,8 +258,8 @@ def test_group_coco_annotations_by_image_id(
         ),  # single image annotations
         (
             [
-                generate_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100),
-                generate_cock_coco_annotation(category_id=0, bbox=(100, 100, 100, 100), area=100 * 100),
+                mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100),
+                mock_cock_coco_annotation(category_id=0, bbox=(100, 100, 100, 100), area=100 * 100),
             ],
             (1000, 1000),
             False,

test/dataset/test_core.py

@@ -0,0 +1,176 @@
+from typing import List, Optional
+
+import pytest
+
+from supervision import DetectionDataset
+from contextlib import ExitStack as DoesNotRaise
+
+import numpy as np
+
+from test.utils import mock_detections
+
+
+@pytest.mark.parametrize(
+    "dataset_list, expected_result, exception",
+    [
+        (
+            [],
+            DetectionDataset(classes=[], images={}, annotations={}),
+            DoesNotRaise()
+        ),  # empty dataset list
+        (
+            [
+                DetectionDataset(classes=[], images={}, annotations={})
+            ],
+            DetectionDataset(classes=[], images={}, annotations={}),
+            DoesNotRaise()
+        ),  # single empty dataset
+        (
+            [
+                DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
+                DetectionDataset(classes=['dog', 'person'], images={}, annotations={})
+            ],
+            DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
+            DoesNotRaise()
+        ),  # two datasets; no images and annotations, the same classes
+        (
+            [
+                DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
+                DetectionDataset(classes=['cat'], images={}, annotations={})
+            ],
+            DetectionDataset(classes=['cat', 'dog', 'person'], images={}, annotations={}),
+            DoesNotRaise()
+        ),  # two datasets; no images and annotations, different classes
+        (
+            [
+                DetectionDataset(
+                    classes=['dog', 'person'],
+                    images={
+                        'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                        'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                        'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    }
+                ),
+                DetectionDataset(classes=[], images={}, annotations={}),
+            ],
+            DetectionDataset(
+                classes=['dog', 'person'],
+                images={
+                    'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                },
+                annotations={
+                    'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                    'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                }
+            ),
+            DoesNotRaise()
+        ),  # two datasets; images and annotations, the same classes
+        (
+            [
+                DetectionDataset(
+                    classes=['dog', 'person'],
+                    images={
+                        'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                        'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                        'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    }
+                ),
+                DetectionDataset(classes=['cat'], images={}, annotations={}),
+            ],
+            DetectionDataset(
+                classes=['cat', 'dog', 'person'],
+                images={
+                    'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                },
+                annotations={
+                    'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
+                }
+            ),
+            DoesNotRaise()
+        ),  # two datasets; images and annotations, different classes
+        (
+            [
+                DetectionDataset(
+                    classes=['dog', 'person'],
+                    images={
+                        'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                        'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                        'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    }
+                ),
+                DetectionDataset(
+                    classes=['cat'],
+                    images={
+                        'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                    }
+                ),
+            ],
+            DetectionDataset(
+                classes=['cat', 'dog', 'person'],
+                images={
+                    'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                },
+                annotations={
+                    'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
+                    'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                }
+            ),
+            DoesNotRaise()
+        ),  # two datasets; images and annotations, different classes
+        (
+            [
+                DetectionDataset(
+                    classes=['dog', 'person'],
+                    images={
+                        'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                        'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                        'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    }
+                ),
+                DetectionDataset(
+                    classes=['dog', 'person'],
+                    images={
+                        'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                        'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
+                    },
+                    annotations={
+                        'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+                        'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+                    }
+                ),
+            ],
+            None,
+            pytest.raises(ValueError)
+        )
+    ]
+)
+def test_dataset_merge(
+    dataset_list: List[DetectionDataset],
+    expected_result: Optional[DetectionDataset],
+    exception: Exception
+) -> None:
+    with exception:
+        result = DetectionDataset.merge(dataset_list=dataset_list)
+        print(result.images.keys())
+        assert result == expected_result

test/dataset/test_utils.py

@@ -1,9 +1,12 @@
-from typing import List, TypeVar, Optional, Tuple
 from contextlib import ExitStack as DoesNotRaise
+from typing import List, TypeVar, Optional, Tuple, Dict
 
 import pytest
 
-from supervision.dataset.utils import train_test_split
+from supervision import Detections
+from supervision.dataset.utils import train_test_split, merge_class_lists, build_class_index_mapping, \
+    map_detections_class_id
+from test.utils import mock_detections
 
 T = TypeVar("T")
 
@@ -80,3 +83,159 @@ def test_train_test_split(
     with exception:
         result = train_test_split(data=data, train_ratio=train_ratio, random_state=random_state, shuffle=shuffle)
         assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'class_lists, expected_result, exception',
+    [
+        (
+            [],
+            [],
+            DoesNotRaise()
+        ),  # empty class lists
+        (
+            [
+                ['dog', 'person']
+            ],
+            ['dog', 'person'],
+            DoesNotRaise()
+        ),  # single class list; already alphabetically sorted
+        (
+            [
+                ['person', 'dog']
+            ],
+            ['dog', 'person'],
+            DoesNotRaise()
+        ),  # single class list; not alphabetically sorted
+        (
+            [
+                ['dog', 'person'],
+                ['dog', 'person']
+            ],
+            ['dog', 'person'],
+            DoesNotRaise()
+        ),  # two class lists; the same classes; already alphabetically sorted
+        (
+            [
+                ['dog', 'person'],
+                ['cat']
+            ],
+            ['cat', 'dog', 'person'],
+            DoesNotRaise()
+        ),  # two class lists; different classes; already alphabetically sorted
+    ]
+)
+def test_merge_class_maps(class_lists: List[List[str]], expected_result: List[str], exception: Exception) -> None:
+    with exception:
+        result = merge_class_lists(class_lists=class_lists)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'source_classes, target_classes, expected_result, exception',
+    [
+        (
+            [],
+            [],
+            {},
+            DoesNotRaise()
+        ),  # empty class lists
+        (
+            [],
+            ['dog', 'person'],
+            {},
+            DoesNotRaise()
+        ),  # empty source class list
+        (
+            ['dog', 'person'],
+            [],
+            None,
+            pytest.raises(ValueError)
+        ),  # empty target class list
+        (
+            ['dog', 'person'],
+            ['dog', 'person'],
+            {0: 0, 1: 1},
+            DoesNotRaise()
+        ),  # same class lists
+        (
+            ['dog', 'person'],
+            ['person', 'dog'],
+            {0: 1, 1: 0},
+            DoesNotRaise()
+        ),  # same class lists but not alphabetically sorted
+        (
+            ['dog', 'person'],
+            ['cat', 'dog', 'person'],
+            {0: 1, 1: 2},
+            DoesNotRaise()
+        ),  # source class list is a subset of target class list
+        (
+            ['dog', 'person'],
+            ['cat', 'dog'],
+            None,
+            pytest.raises(ValueError)
+        ),  # source class list is not a subset of target class list
+    ]
+)
+def test_build_class_index_mapping(
+    source_classes: List[str],
+    target_classes: List[str],
+    expected_result: Optional[Dict[int, int]],
+    exception: Exception
+) -> None:
+    with exception:
+        result = build_class_index_mapping(source_classes=source_classes, target_classes=target_classes)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'source_to_target_mapping, detections, expected_result, exception',
+    [
+        (
+            {},
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+            None,
+            pytest.raises(ValueError)
+        ),  # empty mapping
+        (
+            {0: 1},
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+            DoesNotRaise()
+        ),  # single mapping
+        (
+            {0: 1, 1: 2},
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
+            DoesNotRaise()
+        ),  # multiple mappings
+        (
+            {0: 1, 1: 2},
+            mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[0, 1]),
+            mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[1, 2]),
+            DoesNotRaise()
+        ),  # multiple mappings
+        (
+            {0: 1, 1: 2},
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
+            None,
+            pytest.raises(ValueError)
+        ),  # class_id not in mapping
+        (
+            {0: 1, 1: 2},
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1], confidence=[0.5]),
+            DoesNotRaise(),
+        ),  # confidence is not None
+    ]
+)
+def test_map_detections_class_id(
+    source_to_target_mapping: Dict[int, int],
+    detections: Detections,
+    expected_result: Optional[Detections],
+    exception: Exception
+) -> None:
+    with exception:
+        result = map_detections_class_id(source_to_target_mapping=source_to_target_mapping, detections=detections)
+        assert result == expected_result

test/detection/test_core.py

@@ -8,6 +8,7 @@ from typing import Optional, Union, List
 
 import numpy as np
 
+from test.utils import mock_detections
 
 PREDICTIONS = np.array([
         [       2254,         906,        2447,        1353,     0.90538,           0],
@@ -28,53 +29,30 @@ DETECTIONS = Detections(
 )
 
 
-def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections:
-    return Detections(
-        xyxy = np.array(xyxy, dtype=np.float32),
-        confidence = confidence if confidence is None else np.array(confidence, dtype=np.float32),
-        class_id = class_id if class_id is None else np.array(class_id, dtype=int),
-        tracker_id = tracker_id if tracker_id is None else np.array(tracker_id, dtype=int),
-    )
-
-
 @pytest.mark.parametrize(
     'detections, index, expected_result, exception',
     [
         (
             DETECTIONS,
             DETECTIONS.class_id == 0,
-            Detections(
-                xyxy=np.array([
-                    [       2254,         906,        2447,        1353]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.90538
-                ], dtype=np.float32),
-                class_id=np.array([
-                    0
-                ], dtype=int)
+            mock_detections(
+                xyxy=[[2254, 906, 2447, 1353]],
+                confidence=[0.90538],
+                class_id=[0]
             ),
             DoesNotRaise()
         ),  # take only detections with class_id = 0
         (
             DETECTIONS,
             DETECTIONS.confidence > 0.5,
-            Detections(
-                xyxy=np.array([
-                    [       2254,         906,        2447,        1353],
-                    [       2049,        1133,        2226,        1371],
-                    [        727,        1224,         838,        1601]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.90538,
-                    0.59002,
-                    0.51119
-                ], dtype=np.float32),
-                class_id=np.array([
-                    0,
-                    56,
-                    39
-                ], dtype=int)
+            mock_detections(
+                xyxy=[
+                    [2254, 906, 2447, 1353],
+                    [2049, 1133, 2226, 1371],
+                    [727, 1224, 838, 1601]
+                ],
+                confidence=[0.90538, 0.59002, 0.51119],
+                class_id=[0, 56, 39]
             ),
             DoesNotRaise()
         ),  # take only detections with confidence > 0.5
@@ -97,73 +75,49 @@ def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None)
         (
             DETECTIONS,
             [0, 2],
-            Detections(
-                xyxy=np.array([
-                    [       2254,         906,        2447,        1353],
-                    [        727,        1224,         838,        1601]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.90538,
-                    0.51119
-                ], dtype=np.float32),
-                class_id=np.array([
-                    0,
-                    39
-                ], dtype=int)
+            mock_detections(
+                xyxy=[
+                    [2254, 906, 2447, 1353],
+                    [727, 1224, 838, 1601]
+                ],
+                confidence=[0.90538, 0.51119],
+                class_id=[0, 39]
             ),
             DoesNotRaise()
         ),  # take only first and third detection using List[int] index
         (
             DETECTIONS,
             np.array([0, 2]),
-            Detections(
-                xyxy=np.array([
-                    [       2254,         906,        2447,        1353],
-                    [        727,        1224,         838,        1601]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.90538,
-                    0.51119
-                ], dtype=np.float32),
-                class_id=np.array([
-                    0,
-                    39
-                ], dtype=int)
+            mock_detections(
+                xyxy=[
+                    [2254, 906, 2447, 1353],
+                    [727, 1224, 838, 1601]
+                ],
+                confidence=[0.90538, 0.51119],
+                class_id=[0, 39]
             ),
             DoesNotRaise()
         ),  # take only first and third detection using np.ndarray index
         (
             DETECTIONS,
             0,
-            Detections(
-                xyxy=np.array([
-                    [       2254,         906,        2447,        1353]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.90538
-                ], dtype=np.float32),
-                class_id=np.array([
-                    0
-                ], dtype=int)
+            mock_detections(
+                xyxy=[[2254, 906, 2447, 1353]],
+                confidence=[0.90538],
+                class_id=[0]
             ),
             DoesNotRaise()
         ),  # take only first detection by index
         (
             DETECTIONS,
             slice(1, 3),
-            Detections(
-                xyxy=np.array([
-                    [       2049,        1133,        2226,        1371],
-                    [        727,        1224,         838,        1601]
-                ], dtype=np.float32),
-                confidence=np.array([
-                    0.59002,
-                    0.51119
-                ], dtype=np.float32),
-                class_id=np.array([
-                    56,
-                    39
-                ], dtype=int)
+            mock_detections(
+                xyxy=[
+                    [2049, 1133, 2226, 1371],
+                    [727, 1224, 838, 1601]
+                ],
+                confidence=[0.59002, 0.51119],
+                class_id=[56, 39]
             ),
             DoesNotRaise()
         ),  # take only first detection by index slice (1, 3)

test/utils.py

@@ -0,0 +1,19 @@
+from typing import List
+
+import numpy as np
+
+from supervision import Detections
+
+
+def mock_detections(
+    xyxy: List[List[float]],
+    confidence: List[float] = None,
+    class_id: List[int] = None,
+    tracker_id: List[int] = None
+) -> Detections:
+    return Detections(
+        xyxy=np.array(xyxy, dtype=np.float32),
+        confidence=confidence if confidence is None else np.array(confidence, dtype=np.float32),
+        class_id=class_id if class_id is None else np.array(class_id, dtype=int),
+        tracker_id=tracker_id if tracker_id is None else np.array(tracker_id, dtype=int)
+    )