diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 54fe3614..9931d755 100644 --- a/supervision/dataset/core.py +++ b/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): diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 095e4fde..0a2b90f2 100644 --- a/supervision/dataset/formats/coco.py +++ b/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" ] diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py index 4e38d566..c72ced6f 100644 --- a/supervision/dataset/utils.py +++ b/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: diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index 3b8d8856..ca1ceb6c 100644 --- a/test/dataset/formats/test_coco.py +++ b/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, diff --git a/test/dataset/test_core.py b/test/dataset/test_core.py new file mode 100644 index 00000000..a6252afc --- /dev/null +++ b/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 diff --git a/test/dataset/test_utils.py b/test/dataset/test_utils.py index cf4b93f5..589a8f38 100644 --- a/test/dataset/test_utils.py +++ b/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 diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 0ced48a0..08b6a9d4 100644 --- a/test/detection/test_core.py +++ b/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) diff --git a/test/utils.py b/test/utils.py new file mode 100644 index 00000000..60ae7763 --- /dev/null +++ b/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) + )