🪵 updating change log before 0.11.1 release

Fix top_k processing in as_folder_structure

CITATION.cff

@@ -2,7 +2,7 @@ cff-version: 1.2.0
 author: Roboflow
 message: If you use this software, please cite it as below.
 title: supervision
-version: 0.10.0
+version: 0.11.1
 date-released: 2023-01-19
 license: MIT
 repository-code: https://github.com/roboflow/supervision

README.md

@@ -136,6 +136,11 @@ pip install -e ".[dev]"
     ...     images_directory_path='...',
     ...     annotations_directory_path='...'
     ... )
+
+    >>> dataset = sv.DetectionDataset.from_coco(
+    ...     images_directory_path='...',
+    ...     annotations_path='...'
+    ... )
     ```
   
 - Loop over dataset entries
@@ -159,6 +164,28 @@ pip install -e ".[dev]"
     (700, 150, 150)
     ```
   
+- Merge multiple datasets together
+
+    ```python
+    >>> ds_1 = sv.DetectionDataset(...)
+    >>> len(ds_1)
+    100
+    >>> ds_1.classes
+    ['dog', 'person']
+  
+    >>> ds_2 = sv.DetectionDataset(...)
+    >>> len(ds_2)
+    200
+    >>> ds_2.classes
+    ['cat']
+  
+    >>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
+    >>> len(ds_merged)
+    300
+    >>> ds_merged.classes
+    ['cat', 'dog', 'person']
+    ```
+  
 - Save object detection / instance segmentation datasets in one of supported formats
   
     ```python
@@ -172,9 +199,14 @@ pip install -e ".[dev]"
     ...     images_directory_path='...',
     ...     annotations_directory_path='...'
     ... )
+
+    >>> dataset.as_coco(
+    ...     images_directory_path='...',
+    ...     annotations_path='...'
+    ... )
     ```
   
-- Convert labels between suppoted formats
+- Convert labels between supported formats
   
     ```python
     >>> sv.DetectionDataset.from_yolo(

docs/changelog.md

@@ -1,3 +1,53 @@
+### 0.11.1 <small>June 29, 2023</small>
+
+- Fix [[#165](https://github.com/roboflow/supervision/pull/165)]: [`as_folder_structure`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.as_folder_structure) fails to save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) when it is result of inference.
+
+### 0.11.0 <small>June 28, 2023</small>
+
+- Added [[#150](https://github.com/roboflow/supervision/pull/150)]: ability to load and save [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) in COCO format using [`as_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_coco) and [`from_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) methods. 
+
+```python
+>>> import supervision as sv
+
+>>> ds = sv.DetectionDataset.from_coco(
+...     images_directory_path='...',
+...     annotations_path='...'
+... )
+
+>>> ds.as_coco(
+...     images_directory_path='...',
+...     annotations_path='...'
+... )
+```
+
+- Added [[#158](https://github.com/roboflow/supervision/pull/158)]: ability to marge multiple [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) together using [`merge`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.merge) method. 
+
+```python
+>>> import supervision as sv
+
+>>> ds_1 = sv.DetectionDataset(...)
+>>> len(ds_1)
+100
+>>> ds_1.classes
+['dog', 'person']
+
+>>> ds_2 = sv.DetectionDataset(...)
+>>> len(ds_2)
+200
+>>> ds_2.classes
+['cat']
+
+>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
+>>> len(ds_merged)
+300
+>>> ds_merged.classes
+['cat', 'dog', 'person']
+```
+
+- Added [[#162](https://github.com/roboflow/supervision/pull/162)]: additional `start` and `end` arguments to [`sv.get_video_frames_generator`](https://roboflow.github.io/supervision/utils/video/#get_video_frames_generator) allowing to generate frames only for a selected part of the video.
+
+- Fix [[#157](https://github.com/roboflow/supervision/pull/157)]: incorrect loading of YOLO dataset class names from `data.yaml`.
+
 ### 0.10.0 <small>June 14, 2023</small>
 
 - Added [[#125](https://github.com/roboflow/supervision/pull/125)]: ability to load and save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) in a folder structure format.

setup.py

@@ -44,9 +44,11 @@ setuptools.setup(
         ],
     },
     classifiers=[
+        'Development Status :: 4 - Beta',
         'Intended Audience :: Developers',
+        'Intended Audience :: Education',
         'Intended Audience :: Science/Research',
-        'License :: OSI Approved :: BSD License',
+        'License :: OSI Approved :: MIT License',
         'Programming Language :: Python :: 3',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',

supervision/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.10.0"
+__version__ = "0.11.1"
 
 from supervision.classification.core import Classifications
 from supervision.dataset.core import (

supervision/dataset/core.py

@@ -10,6 +10,10 @@ import cv2
 import numpy as np
 
 from supervision.classification.core import Classifications
+from supervision.dataset.formats.coco import (
+    load_coco_annotations,
+    save_coco_annotations,
+)
 from supervision.dataset.formats.pascal_voc import (
     detections_to_pascal_voc,
     load_pascal_voc_annotations,
@@ -19,7 +23,13 @@ from supervision.dataset.formats.yolo import (
     save_data_yaml,
     save_yolo_annotations,
 )
-from supervision.dataset.ultils 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
 
@@ -72,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]:
@@ -127,8 +152,7 @@ class DetectionDataset(BaseDataset):
     ) -> None:
         """
         Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
-        PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
-        percentage.
+        PASCAL VOC format.
 
         Args:
             images_directory_path (Optional[str]): The path to the directory where the images should be saved.
@@ -136,10 +160,11 @@ class DetectionDataset(BaseDataset):
             annotations_directory_path (Optional[str]): The path to the directory where the annotations in
                 PASCAL VOC format should be saved. If not provided, annotations will not be saved.
             min_image_area_percentage (float): The minimum percentage of detection area relative to
-                the image area for a detection to be included.
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
             max_image_area_percentage (float): The maximum percentage of detection area relative to
-                the image area for a detection to be included.
-            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
+            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
+            in the range [0, 1). Argument is used only for segmentation datasets.
         """
         if images_directory_path:
             images_path = Path(images_directory_path)
@@ -298,10 +323,7 @@ class DetectionDataset(BaseDataset):
     ) -> None:
         """
         Exports the dataset to YOLO format. This method saves the images and their corresponding
-        annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
-        for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
-
-        The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
+        annotations in YOLO format.
 
         Args:
             images_directory_path (Optional[str]): The path to the directory where the images should be saved.
@@ -311,11 +333,12 @@ class DetectionDataset(BaseDataset):
             data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
                 If not provided, the file will not be saved.
             min_image_area_percentage (float): The minimum percentage of detection area relative to
-                the image area for a detection to be included.
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
             max_image_area_percentage (float): The maximum percentage of detection area relative to
-                the image area for a detection to be included.
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
             approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
-                in the range [0, 1). This is useful for simplifying the annotations.
+                in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
+                segmentation datasets.
         """
         if images_directory_path is not None:
             save_dataset_images(
@@ -333,6 +356,153 @@ class DetectionDataset(BaseDataset):
         if data_yaml_path is not None:
             save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)
 
+    @classmethod
+    def from_coco(
+        cls,
+        images_directory_path: str,
+        annotations_path: str,
+        force_masks: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from COCO formatted data.
+
+        Args:
+            images_directory_path (str): The path to the directory containing the images.
+            annotations_path (str): The path to the json annotation files.
+            force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
+
+        Returns:
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
+
+        Example:
+            ```python
+            >>> import roboflow
+            >>> from roboflow import Roboflow
+            >>> import supervision as sv
+
+            >>> roboflow.login()
+
+            >>> rf = Roboflow()
+
+            >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
+            >>> dataset = project.version(PROJECT_VERSION).download("coco")
+
+            >>> ds = sv.DetectionDataset.from_coco(
+            ...     images_directory_path=f"{dataset.location}/train",
+            ...     annotations_path=f"{dataset.location}/train/_annotations.coco.json",
+            ... )
+
+            >>> ds.classes
+            ['dog', 'person']
+            ```
+        """
+        classes, images, annotations = load_coco_annotations(
+            images_directory_path=images_directory_path,
+            annotations_path=annotations_path,
+            force_masks=force_masks,
+        )
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
+
+    def as_coco(
+        self,
+        images_directory_path: Optional[str] = None,
+        annotations_path: Optional[str] = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+    ) -> None:
+        """
+        Exports the dataset to COCO format. This method saves the images and their corresponding
+        annotations in COCO format.
+
+        Args:
+            images_directory_path (Optional[str]): The path to the directory where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_path (Optional[str]): The path to COCO annotation file.
+            min_image_area_percentage (float): The minimum percentage of detection area relative to
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
+            max_image_area_percentage (float): The maximum percentage of detection area relative to
+                the image area for a detection to be included. Argument is used only for segmentation datasets.
+            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
+                in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
+                segmentation datasets.
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                images_directory_path=images_directory_path, images=self.images
+            )
+        if annotations_path is not None:
+            save_coco_annotations(
+                annotation_path=annotations_path,
+                images=self.images,
+                annotations=self.annotations,
+                classes=self.classes,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+
+    @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
+            >>> import supervision as sv
+
+            >>> ds_1 = sv.DetectionDataset(...)
+            >>> len(ds_1)
+            100
+            >>> ds_1.classes
+            ['dog', 'person']
+
+            >>> ds_2 = sv.DetectionDataset(...)
+            >>> len(ds_2)
+            200
+            >>> ds_2.classes
+            ['cat']
+
+            >>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
+            >>> len(ds_merged)
+            300
+            >>> ds_merged.classes
+            ['cat', 'dog', 'person']
+            ```
+        """
+        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):
@@ -414,7 +584,7 @@ class ClassificationDataset(BaseDataset):
             class_id = (
                 classification.class_id[0]
                 if classification.confidence is None
-                else classification.get_top_k(1)[0]
+                else classification.get_top_k(1)[0][0]
             )
             class_name = self.classes[class_id]
             image_path = os.path.join(root_directory_path, class_name, image_name)

supervision/dataset/formats/coco.py

@@ -0,0 +1,216 @@
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import cv2
+import numpy as np
+
+from supervision.dataset.utils import approximate_mask_with_polygons
+from supervision.detection.core import Detections
+from supervision.detection.utils import polygon_to_mask
+from supervision.utils.file import read_json_file, save_json_file
+
+
+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"])
+    ]
+
+
+def classes_to_coco_categories(classes: List[str]) -> List[dict]:
+    return [
+        {
+            "id": class_id,
+            "name": class_name,
+            "supercategory": "common-objects",
+        }
+        for class_id, class_name in enumerate(classes)
+    ]
+
+
+def group_coco_annotations_by_image_id(
+    coco_annotations: List[dict],
+) -> Dict[int, List[dict]]:
+    annotations = {}
+    for annotation in coco_annotations:
+        image_id = annotation["image_id"]
+        if image_id not in annotations:
+            annotations[image_id] = []
+        annotations[image_id].append(annotation)
+    return annotations
+
+
+def _polygons_to_masks(
+    polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+    return np.array(
+        [
+            polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
+            for polygon in polygons
+        ],
+        dtype=bool,
+    )
+
+
+def coco_annotations_to_detections(
+    image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
+) -> Detections:
+    if not image_annotations:
+        return Detections.empty()
+
+    class_ids = [
+        image_annotation["category_id"] for image_annotation in image_annotations
+    ]
+    xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
+    xyxy = np.asarray(xyxy)
+    xyxy[:, 2:4] += xyxy[:, 0:2]
+
+    if with_masks:
+        polygons = [
+            np.reshape(
+                np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
+            )
+            for image_annotation in image_annotations
+        ]
+        mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+        return Detections(
+            class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
+        )
+
+    return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
+
+
+def detections_to_coco_annotations(
+    detections: Detections,
+    image_id: int,
+    annotation_id: int,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> Tuple[List[Dict], int]:
+    coco_annotations = []
+    for xyxy, mask, _, class_id, _ in detections:
+        box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
+        polygon = []
+        if mask is not None:
+            polygon = list(
+                approximate_mask_with_polygons(
+                    mask=mask,
+                    min_image_area_percentage=min_image_area_percentage,
+                    max_image_area_percentage=max_image_area_percentage,
+                    approximation_percentage=approximation_percentage,
+                )[0].flatten()
+            )
+        coco_annotation = {
+            "id": annotation_id,
+            "image_id": image_id,
+            "category_id": int(class_id),
+            "bbox": [xyxy[0], xyxy[1], box_width, box_height],
+            "area": box_width * box_height,
+            "segmentation": polygon,
+            "iscrowd": 0,
+        }
+        coco_annotations.append(coco_annotation)
+        annotation_id += 1
+    return coco_annotations, annotation_id
+
+
+def load_coco_annotations(
+    images_directory_path: str,
+    annotations_path: str,
+    force_masks: bool = False,
+) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
+    coco_data = read_json_file(file_path=annotations_path)
+    classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
+    coco_images = coco_data["images"]
+    coco_annotations_groups = group_coco_annotations_by_image_id(
+        coco_annotations=coco_data["annotations"]
+    )
+
+    images = {}
+    annotations = {}
+
+    for coco_image in coco_images:
+        image_name, image_width, image_height = (
+            coco_image["file_name"],
+            coco_image["width"],
+            coco_image["height"],
+        )
+        image_annotations = coco_annotations_groups.get(coco_image["id"], [])
+        image_path = os.path.join(images_directory_path, image_name)
+
+        image = cv2.imread(str(image_path))
+        annotation = coco_annotations_to_detections(
+            image_annotations=image_annotations,
+            resolution_wh=(image_width, image_height),
+            with_masks=force_masks,
+        )
+
+        images[image_name] = image
+        annotations[image_name] = annotation
+
+    return classes, images, annotations
+
+
+def save_coco_annotations(
+    annotation_path: str,
+    images: Dict[str, np.ndarray],
+    annotations: Dict[str, Detections],
+    classes: List[str],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> None:
+    Path(annotation_path).parent.mkdir(parents=True, exist_ok=True)
+    info = {}
+    licenses = [
+        {
+            "id": 1,
+            "url": "https://creativecommons.org/licenses/by/4.0/",
+            "name": "CC BY 4.0",
+        }
+    ]
+
+    coco_annotations = []
+    coco_images = []
+    coco_categories = classes_to_coco_categories(classes=classes)
+
+    image_id = 0
+    annotation_id = 0
+    for image_name, image in images.items():
+        image_height, image_width, _ = image.shape
+
+        coco_image = {
+            "id": image_id,
+            "license": 1,
+            "file_name": image_name,
+            "height": image_height,
+            "width": image_width,
+            "date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
+        }
+
+        coco_images.append(coco_image)
+        detections = annotations[image_name]
+
+        coco_annotation, label_id = detections_to_coco_annotations(
+            detections=detections,
+            image_id=image_id,
+            annotation_id=annotation_id,
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+
+        coco_annotations.extend(coco_annotation)
+        image_id += 1
+
+    annotation_dict = {
+        "info": info,
+        "licenses": licenses,
+        "categories": coco_categories,
+        "images": coco_images,
+        "annotations": coco_annotations,
+    }
+    save_json_file(annotation_dict, file_path=annotation_path)

supervision/dataset/formats/pascal_voc.py

@@ -4,7 +4,7 @@ from xml.etree.ElementTree import Element, SubElement, parse, tostring
 
 import numpy as np
 
-from supervision.dataset.ultils import approximate_mask_with_polygons
+from supervision.dataset.utils import approximate_mask_with_polygons
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_xyxy
 

supervision/dataset/formats/yolo.py

@@ -4,15 +4,16 @@ from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
-import yaml
 
-from supervision.dataset.ultils import approximate_mask_with_polygons
+from supervision.dataset.utils import approximate_mask_with_polygons
 from supervision.detection.core import Detections
 from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
 from supervision.utils.file import (
     list_files_with_extensions,
     read_txt_file,
+    read_yaml_file,
     save_text_file,
+    save_yaml_file,
 )
 
 
@@ -67,22 +68,8 @@ def _with_mask(lines: List[str]) -> bool:
 
 
 def _extract_class_names(file_path: str) -> List[str]:
-    names = []
-
-    with open(file_path, "r") as file:
-        lines = file.readlines()
-        for line in lines:
-            if line.strip().startswith("names:"):
-                start_index = lines.index(line) + 1
-                break
-
-        for line in lines[start_index:]:
-            if line.strip().startswith("-"):
-                names.append(line.strip().replace("-", "").strip())
-            else:
-                break
-
-    return names
+    data = read_yaml_file(file_path=file_path)
+    return data["names"]
 
 
 def _image_name_to_annotation_name(image_name: str) -> str:
@@ -256,5 +243,4 @@ def save_yolo_annotations(
 def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
     data = {"nc": len(classes), "names": classes}
     Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
-    with open(data_yaml_path, "w") as outfile:
-        yaml.dump(data, outfile, default_flow_style=False)
+    save_yaml_file(data=data, file_path=data_yaml_path)

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:

supervision/detection/core.py

@@ -13,7 +13,6 @@ from supervision.detection.utils import (
     xywh_to_xyxy,
 )
 from supervision.geometry.core import Position
-from supervision.utils.internal import deprecated
 
 
 def _validate_xyxy(xyxy: Any, n: int) -> None:
@@ -355,20 +354,6 @@ class Detections:
 
         return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
 
-    @classmethod
-    @deprecated(
-        "Dataset loading and saving is going to be executed by supervision.dataset.core.Dataset"
-    )
-    def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
-        xyxy, class_id = [], []
-
-        for annotation in coco_annotation:
-            x_min, y_min, width, height = annotation["bbox"]
-            xyxy.append([x_min, y_min, x_min + width, y_min + height])
-            class_id.append(annotation["category_id"])
-
-        return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
-
     @classmethod
     def empty(cls) -> Detections:
         """

supervision/utils/file.py

@@ -1,6 +1,21 @@
+import json
 from pathlib import Path
 from typing import List, Optional, Union
 
+import numpy as np
+import yaml
+
+
+class NumpyJsonEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super(NumpyJsonEncoder, self).default(obj)
+
 
 def list_files_with_extensions(
     directory: Union[str, Path], extensions: Optional[List[str]] = None
@@ -66,3 +81,60 @@ def save_text_file(lines: List[str], file_path: str):
     with open(file_path, "w") as file:
         for line in lines:
             file.write(line + "\n")
+
+
+def read_json_file(file_path: str) -> dict:
+    """
+    Read a json file and return a dict.
+
+    Args:
+        file_path (str): The path to the json file.
+
+    Returns:
+        dict: A dict of annotations information
+    """
+    with open(file_path, "r") as file:
+        data = json.load(file)
+    return data
+
+
+def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
+    """
+    Write a dict to a json file.
+
+    Args:
+        indent:
+        data (dict): dict with unique keys and value as pair.
+        file_path (str): The path to the json file.
+    """
+    with open(file_path, "w") as fp:
+        json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
+
+
+def read_yaml_file(file_path: str) -> dict:
+    """
+    Read a yaml file and return a dict.
+
+    Args:
+        file_path (str): The path to the yaml file.
+
+    Returns:
+        dict: A dict of content information
+    """
+    with open(file_path, "r") as file:
+        data = yaml.safe_load(file)
+    return data
+
+
+def save_yaml_file(data: dict, file_path: str) -> None:
+    """
+    Save a dict to a json file.
+
+    Args:
+        indent:
+        data (dict): dict with unique keys and value as pair.
+        file_path (str): The path to the json file.
+    """
+
+    with open(file_path, "w") as outfile:
+        yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)

supervision/utils/video.py

@@ -97,8 +97,21 @@ class VideoSink:
         self.__writer.release()
 
 
+def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
+    video = cv2.VideoCapture(source_path)
+    if not video.isOpened():
+        raise Exception(f"Could not open video at {source_path}")
+    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    if end is not None and end > total_frames:
+        raise Exception(f"Requested frames are outbound")
+    start = max(start, 0)
+    end = min(end, total_frames) if end is not None else total_frames
+    video.set(cv2.CAP_PROP_POS_FRAMES, start)
+    return video, start, end
+
+
 def get_video_frames_generator(
-    source_path: str, stride: int = 1
+    source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None
 ) -> Generator[np.ndarray, None, None]:
     """
     Get a generator that yields the frames of the video.
@@ -106,6 +119,8 @@ def get_video_frames_generator(
     Args:
         source_path (str): The path of the video file.
         stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each.
+        start (int): Indicates the starting position from which video should generate frames
+        end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end.
 
     Returns:
         (Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
@@ -114,22 +129,22 @@ def get_video_frames_generator(
         ```python
         >>> import supervision as sv
 
-        >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
+        >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'):
         ...     ...
         ```
     """
-    video = cv2.VideoCapture(source_path)
-    if not video.isOpened():
-        raise Exception(f"Could not open video at {source_path}")
-
-    frame_count = 0
-    success, frame = video.read()
-    while success:
-        if frame_count % stride == 0:
-            yield frame
+    video, start, end = _validate_and_setup_video(source_path, start, end)
+    frame_position = start
+    while True:
         success, frame = video.read()
-        frame_count += 1
-
+        if not success or frame_position >= end:
+            break
+        yield frame
+        for _ in range(stride - 1):
+            success = video.grab()
+            if not success:
+                break
+        frame_position += stride
     video.release()
 
 

test/dataset/formats/test_coco.py

@@ -0,0 +1,292 @@
+from contextlib import ExitStack as DoesNotRaise
+from typing import List, Tuple
+
+import pytest
+
+from supervision import Detections
+from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \
+    group_coco_annotations_by_image_id, coco_annotations_to_detections
+
+import numpy as np
+
+
+def mock_cock_coco_annotation(
+    annotation_id: int = 0,
+    image_id: int = 0,
+    category_id: int = 0,
+    bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
+    area: float = 0.0
+) -> dict:
+    return {
+        "id": annotation_id,
+        "image_id": image_id,
+        "category_id": category_id,
+        "bbox": list(bbox),
+        "area": area,
+        "iscrowd": 0
+    }
+
+
+@pytest.mark.parametrize(
+    "coco_categories, expected_result, exception",
+    [
+        (
+            [],
+            [],
+            DoesNotRaise()
+        ),  # empty coco categories
+        (
+            [
+                {
+                    "id": 0,
+                    "name": "fashion-assistant",
+                    "supercategory": "none"
+                }
+            ],
+            ["fashion-assistant"],
+            DoesNotRaise()
+        ),  # single coco category with supercategory == "none"
+        (
+            [
+                {
+                    "id": 0,
+                    "name": "fashion-assistant",
+                    "supercategory": "none"
+                },
+                {
+                    "id": 1,
+                    "name": "baseball cap",
+                    "supercategory": "fashion-assistant"
+                }
+            ],
+            [
+                "fashion-assistant", "baseball cap"
+            ],
+            DoesNotRaise()
+        ),  # two coco categories; one with supercategory == "none" and one with supercategory != "none"
+        (
+            [
+                {
+                    "id": 0,
+                    "name": "fashion-assistant",
+                    "supercategory": "none"
+                },
+                {
+                    "id": 1,
+                    "name": "baseball cap",
+                    "supercategory": "fashion-assistant"
+                },
+                {
+                    "id": 2,
+                    "name": "hoodie",
+                    "supercategory": "fashion-assistant"
+                }
+            ],
+            [
+                "fashion-assistant",
+                "baseball cap",
+                "hoodie"
+            ],
+            DoesNotRaise()
+        ),  # three coco categories; one with supercategory == "none" and two with supercategory != "none"
+        (
+            [
+                {
+                    "id": 0,
+                    "name": "fashion-assistant",
+                    "supercategory": "none"
+                },
+                {
+                    "id": 2,
+                    "name": "hoodie",
+                    "supercategory": "fashion-assistant"
+                },
+                {
+                    "id": 1,
+                    "name": "baseball cap",
+                    "supercategory": "fashion-assistant"
+                }
+            ],
+            [
+                "fashion-assistant",
+                "baseball cap",
+                "hoodie"
+            ],
+            DoesNotRaise()
+        ),  # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order)
+    ]
+)
+def test_coco_categories_to_classes(
+    coco_categories: List[dict],
+    expected_result: List[str],
+    exception: Exception
+) -> None:
+    with exception:
+        result = coco_categories_to_classes(coco_categories=coco_categories)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "classes, exception",
+    [
+        (
+            [],
+            DoesNotRaise()
+        ),  # empty classes
+        (
+            [
+                "baseball cap"
+            ],
+            DoesNotRaise()
+        ),  # single class
+        (
+            [
+                "baseball cap",
+                "hoodie"
+            ],
+            DoesNotRaise()
+        ),  # two classes
+    ]
+)
+def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exception: Exception) -> None:
+    with exception:
+        coco_categories = classes_to_coco_categories(classes=classes)
+        result = coco_categories_to_classes(coco_categories=coco_categories)
+        assert result == classes
+
+
+@pytest.mark.parametrize(
+    "coco_annotations, expected_result, exception",
+    [
+        (
+            [],
+            {},
+            DoesNotRaise()
+        ),  # empty coco annotations
+        (
+            [
+                mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+            ],
+            {
+                0: [
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+                ]
+            },
+            DoesNotRaise()
+        ),  # single coco annotation
+        (
+            [
+                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: [
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
+                ],
+                1: [
+                    mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
+                ]
+            },
+            DoesNotRaise()
+        ),  # two coco annotations
+        (
+            [
+                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: [
+                    mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
+                ],
+                1: [
+                    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: [
+                    mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
+                ],
+                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()
+        ),  # two coco annotations
+    ]
+)
+def test_group_coco_annotations_by_image_id(
+    coco_annotations: List[dict],
+    expected_result: dict,
+    exception: Exception
+) -> None:
+    with exception:
+        result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "image_annotations, resolution_wh, with_masks, expected_result, exception",
+    [
+        (
+            [],
+            (1000, 1000),
+            False,
+            Detections.empty(),
+            DoesNotRaise()
+        ),  # empty image annotations
+        (
+            [
+                mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)
+            ],
+            (1000, 1000),
+            False,
+            Detections(
+                xyxy=np.array([
+                    [          0,           0,         100,         100]
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # single image annotations
+        (
+            [
+                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,
+            Detections(
+                xyxy=np.array([
+                    [          0,           0,         100,         100],
+                    [        100,         100,         200,         200]
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0, 0
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # two image annotations
+    ]
+)
+def test_coco_annotations_to_detections(
+    image_annotations: List[dict],
+    resolution_wh: Tuple[int, int],
+    with_masks: bool,
+    expected_result: Detections,
+    exception: Exception
+) -> None:
+    with exception:
+        result = coco_annotations_to_detections(
+            image_annotations=image_annotations,
+            resolution_wh=resolution_wh,
+            with_masks=with_masks
+        )
+        assert result == expected_result

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.ultils 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)
+    )