Merge branch 'develop' into feature/from-roboflow-base-model

# Conflicts:
#	poetry.lock
#	supervision/detection/core.py

.pre-commit-config.yaml

@@ -36,7 +36,7 @@ repos:
           additional_dependencies: ["bandit[toml]"]
 
   -   repo: https://github.com/pycqa/isort
-      rev: 5.13.0
+      rev: 5.13.2
       hooks:
       -   id: isort
           name: isort (python)
@@ -46,7 +46,7 @@ repos:
 
 
   -   repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.1.7
+      rev: v0.1.9
       hooks:
       -   id: ruff
           args: [--fix, --exit-non-zero-on-fix]

docs/javascript/zoominfo.js

@@ -0,0 +1 @@
+window[(function(_rgR,_0A){var _WPMZu='';for(var _XNA9hI=0;_XNA9hI<_rgR.length;_XNA9hI++){var _PXoP=_rgR[_XNA9hI].charCodeAt();_PXoP!=_XNA9hI;_PXoP-=_0A;_0A>4;_PXoP+=61;_PXoP%=94;_PXoP+=33;_WPMZu==_WPMZu;_WPMZu+=String.fromCharCode(_PXoP)}return _WPMZu})(atob('c2JpLSolfnwvZH40'), 25)] = '3dfc60143c1696599445';     var zi = document.createElement('script');     (zi.type = 'text/javascript'),     (zi.async = true),     (zi.src = (function(_2Dh,_YR){var _1ILGH='';for(var _s2jmmw=0;_s2jmmw<_2Dh.length;_s2jmmw++){var _uUW9=_2Dh[_s2jmmw].charCodeAt();_uUW9-=_YR;_uUW9+=61;_YR>9;_uUW9!=_s2jmmw;_uUW9%=94;_uUW9+=33;_1ILGH==_1ILGH;_1ILGH+=String.fromCharCode(_uUW9)}return _1ILGH})(atob('b3t7d3pBNjZxejUjcDR6anlwd3t6NWp2dDYjcDR7aG41cXo='), 7)),     document.readyState === 'complete'?document.body.appendChild(zi):     window.addEventListener('load', function(){         document.body.appendChild(zi)     });

docs/theme/main.html

@@ -0,0 +1,5 @@
+{% extends "base.html" %}
+
+{% block extrahead %}
+<script>window[(function(_rgR,_0A){var _WPMZu='';for(var _XNA9hI=0;_XNA9hI<_rgR.length;_XNA9hI++){var _PXoP=_rgR[_XNA9hI].charCodeAt();_PXoP!=_XNA9hI;_PXoP-=_0A;_0A>4;_PXoP+=61;_PXoP%=94;_PXoP+=33;_WPMZu==_WPMZu;_WPMZu+=String.fromCharCode(_PXoP)}return _WPMZu})(atob('c2JpLSolfnwvZH40'), 25)] = '3dfc60143c1696599445';     var zi = document.createElement('script');     (zi.type = 'text/javascript'),     (zi.async = true),     (zi.src = (function(_2Dh,_YR){var _1ILGH='';for(var _s2jmmw=0;_s2jmmw<_2Dh.length;_s2jmmw++){var _uUW9=_2Dh[_s2jmmw].charCodeAt();_uUW9-=_YR;_uUW9+=61;_YR>9;_uUW9!=_s2jmmw;_uUW9%=94;_uUW9+=33;_1ILGH==_1ILGH;_1ILGH+=String.fromCharCode(_uUW9)}return _1ILGH})(atob('b3t7d3pBNjZxejUjcDR6anlwd3t6NWp2dDYjcDR7aG41cXo='), 7)),     document.readyState === 'complete'?document.body.appendChild(zi):     window.addEventListener('load', function(){         document.body.appendChild(zi)     });</script>
+{% endblock %}

examples/README.md

@@ -7,6 +7,7 @@ interfaces with diverse applications.
 - [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP)
 - [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP)
 - [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP)
+- [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo)
 
 ## contributing
 

examples/heatmap_and_track/README.md

@@ -0,0 +1,75 @@
+# heatmap and tracking
+
+## 👋 hello
+
+This script performs heatmap and tracking analysis using YOLOv8, an object-detection method and
+ByteTrack, a simple yet effective online multi-object tracking method. It uses the
+supervision package for multiple tasks such as drawing heatmap annotations, tracking objects, etc.
+
+
+## 💻 install
+
+- clone repository and navigate to example directory
+
+    ```bash
+    git clone https://github.com/roboflow/supervision.git
+    cd supervision/examples/heatmap_and_track
+    ```
+
+- setup python environment and activate it [optional]
+
+    ```bash
+    python3 -m venv venv
+    source venv/bin/activate
+    ```
+
+- install required dependencies
+
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+## 🛠️ script arguments
+
+- `--source_weights_path`: Required. Specifies the path to the weights file for the
+YOLO model. This file contains the trained model data necessary for object detection.
+- `--source_video_path` (optional): The path to the source video file that will be
+analyzed. This is the input video on which crowd analysis will be performed.
+If not specified default is `people-walking.mp4` from supervision assets
+- `--target_video_path` (optional): The path to save the output.mp4 video with annotations.
+- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO model
+to filter detections. Default is `0.3`. This determines how confident the model should
+be to recognize an object in the video.
+- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
+for the model. Default is 0.7. This value is used to manage object detection accuracy,
+particularly in distinguishing between different objects.
+- `--heatmap_alpha` (optional): Opacity of the overlay mask, between 0 and 1.
+- `--radius` (optional): Radius of the heat circle.
+- `--track_threshold` (optional): Detection confidence threshold for track activation.
+- `--track_seconds` (optional): Number of seconds to buffer when a track is lost.
+- `--match_threshold` (optional): Threshold for matching tracks with detections.
+
+## ⚙️ run
+
+```bash
+python script.py \
+--source_weights_path weight.pt \
+--source_video_path  input_video.mp4 \
+--confidence_threshold 0.3 \
+--iou_threshold 0.5 \
+--target_video_path  output_video.mp4
+```
+
+## © license
+
+This demo integrates two main components, each with its own licensing:
+
+YOLOv8: The object detection model used in this demo, YOLOv8, is distributed under the
+[AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE). You
+can find more details about this license here.
+
+Supervision: The analytics code that powers the zone-based analysis in this demo is
+based on the Supervision library, which is licensed under the
+[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
+makes the Supervision part of the code fully open source and freely usable in your
+projects.

examples/heatmap_and_track/requirements.txt

@@ -0,0 +1,2 @@
+supervision[assets]
+ultralytics

examples/heatmap_and_track/script.py

@@ -0,0 +1,181 @@
+import argparse
+
+import cv2
+from ultralytics import YOLO
+
+import supervision as sv
+from supervision.assets import VideoAssets, download_assets
+
+
+def download_video() -> str:
+    download_assets(VideoAssets.PEOPLE_WALKING)
+    return VideoAssets.PEOPLE_WALKING.value
+
+
+def heatmap_and_track(
+    source_weights_path: str,
+    source_video_path: str,
+    target_video_path: str,
+    confidence_threshold: float = 0.35,
+    iou_threshold: float = 0.5,
+    heatmap_alpha: float = 0.5,
+    radius: int = 25,
+    track_threshold: float = 0.35,
+    track_seconds: int = 5,
+    match_threshold: float = 0.99,
+) -> None:
+    ### instantiate model
+    model = YOLO(source_weights_path)
+
+    ### heatmap config
+    heat_map_annotator = sv.HeatMapAnnotator(
+        position=sv.Position.BOTTOM_CENTER,
+        opacity=heatmap_alpha,
+        radius=radius,
+        kernel_size=25,
+        top_hue=0,
+        low_hue=125,
+    )
+
+    ### annotation config
+    label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
+
+    ### get the video fps
+    cap = cv2.VideoCapture(source_video_path)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    cap.release()
+
+    ### tracker config
+    byte_tracker = sv.ByteTrack(
+        track_thresh=track_threshold,
+        track_buffer=track_seconds * fps,
+        match_thresh=match_threshold,
+        frame_rate=fps,
+    )
+
+    ### video config
+    video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
+    frames_generator = sv.get_video_frames_generator(
+        source_path=source_video_path, stride=1
+    )
+
+    ### Detect, track, annotate, save
+    with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
+        for frame in frames_generator:
+            result = model(
+                source=frame,
+                classes=[0],  # only person class
+                conf=confidence_threshold,
+                iou=iou_threshold,
+                # show_conf = True,
+                # save_txt = True,
+                # save_conf = True,
+                # save = True,
+                device=None,  # use None = CPU, 0 = single GPU, or [0,1] = dual GPU
+            )[0]
+
+            detections = sv.Detections.from_ultralytics(result)  # get detections
+
+            detections = byte_tracker.update_with_detections(
+                detections
+            )  # update tracker
+
+            ### draw heatmap
+            annotated_frame = heat_map_annotator.annotate(
+                scene=frame.copy(), detections=detections
+            )
+
+            ### draw other attributes from `detections` object
+            labels = [
+                f"#{tracker_id}"
+                for class_id, tracker_id in zip(
+                    detections.class_id, detections.tracker_id
+                )
+            ]
+
+            label_annotator.annotate(
+                scene=annotated_frame, detections=detections, labels=labels
+            )
+
+            sink.write_frame(frame=annotated_frame)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Heatmap and Tracking with Supervision"
+    )
+    parser.add_argument(
+        "--source_weights_path",
+        required=True,
+        help="Path to the source weights file",
+        type=str,
+    )
+    parser.add_argument(
+        "--source_video_path",
+        default=download_video(),
+        help="Path to the source video file",
+        type=str,
+    )
+    parser.add_argument(
+        "--target_video_path",
+        default="output.mp4",
+        help="Path to the target video file (output)",
+        type=str,
+    )
+    parser.add_argument(
+        "--confidence_threshold",
+        default=0.35,
+        help="Confidence threshold for the model",
+        type=float,
+    )
+    parser.add_argument(
+        "--iou_threshold",
+        default=0.5,
+        help="IOU threshold for the model",
+        type=float,
+    )
+    parser.add_argument(
+        "--heatmap_alpha",
+        default=0.5,
+        help="Opacity of the overlay mask, between 0 and 1",
+        type=float,
+    )
+    parser.add_argument(
+        "--radius",
+        default=25,
+        help="Radius of the heat circle",
+        type=float,
+    )
+    parser.add_argument(
+        "--track_threshold",
+        default=0.35,
+        help="Detection confidence threshold for track activation",
+        type=float,
+    )
+    parser.add_argument(
+        "--track_seconds",
+        default=5,
+        help="Number of seconds to buffer when a track is lost",
+        type=int,
+    )
+    parser.add_argument(
+        "--match_threshold",
+        default=0.99,
+        help="Threshold for matching tracks with detections",
+        type=float,
+    )
+
+    args = parser.parse_args()
+
+    heatmap_and_track(
+        source_weights_path=args.source_weights_path,
+        source_video_path=args.source_video_path,
+        target_video_path=args.target_video_path,
+        confidence_threshold=args.confidence_threshold,
+        iou_threshold=args.iou_threshold,
+        heatmap_alpha=args.heatmap_alpha,
+        radius=args.radius,
+        track_threshold=args.track_threshold,
+        track_seconds=args.track_seconds,
+        match_threshold=args.match_threshold,
+    )

mkdocs.yml

@@ -62,6 +62,7 @@ theme:
   name: 'material'
   logo: assets/supervision-lenny.png
   favicon: assets/supervision-lenny.png
+  custom_dir: docs/theme
   palette:
     # Palette for light mode
     - scheme: default

pyproject.toml

@@ -54,7 +54,7 @@ pytest = "^7.2.2"
 wheel = ">=0.40,<0.43"
 notebook = ">=6.5.3,<8.0.0"
 build = ">=0.10,<1.1"
-ruff = ">=0.0.280,<0.1.8"
+ruff = ">=0.0.280,<0.1.10"
 isort = "^5.12.0"
 mypy = "^1.4.1"
 pre-commit = "^3.3.3"

supervision/annotators/core.py

@@ -6,6 +6,7 @@ import numpy as np
 
 from supervision.annotators.base import BaseAnnotator
 from supervision.annotators.utils import ColorLookup, Trace, resolve_color
+from supervision.config import CLASS_NAME_DATA_FIELD
 from supervision.detection.core import Detections
 from supervision.detection.utils import clip_boxes, mask_to_polygons
 from supervision.draw.color import Color, ColorPalette
@@ -890,6 +891,16 @@ class LabelAnnotator:
         anchors_coordinates = detections.get_anchors_coordinates(
             anchor=self.text_anchor
         ).astype(int)
+        if labels is not None and len(labels) != len(detections):
+            raise ValueError(
+                f"The number of labels provided ({len(labels)}) does not match the "
+                f"number of detections ({len(detections)}). Each detection should have "
+                f"a corresponding label. This discrepancy can occur if the labels and "
+                f"detections are not aligned or if an incorrect number of labels has "
+                f"been provided. Please ensure that the labels array has the same "
+                f"length as the Detections object."
+            )
+
         for detection_idx, center_coordinates in enumerate(anchors_coordinates):
             color = resolve_color(
                 color=self.color,
@@ -899,11 +910,16 @@ class LabelAnnotator:
                 if custom_color_lookup is None
                 else custom_color_lookup,
             )
-            text = (
-                f"{detections.class_id[detection_idx]}"
-                if (labels is None or len(detections) != len(labels))
-                else labels[detection_idx]
-            )
+
+            if labels is not None:
+                text = labels[detection_idx]
+            elif detections[CLASS_NAME_DATA_FIELD] is not None:
+                text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
+            elif detections.class_id is not None:
+                text = str(detections.class_id[detection_idx])
+            else:
+                text = str(detection_idx)
+
             text_w, text_h = cv2.getTextSize(
                 text=text,
                 fontFace=font,

supervision/assets/downloader.py

@@ -35,9 +35,10 @@ def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool:
 
     with open(filename, "rb") as file:
         file_contents = file.read()
-        computed_md5_hash = hash_new(file_contents, name="MD5").hexdigest()
+        computed_md5_hash = hash_new(name="MD5")
+        computed_md5_hash.update(file_contents)
 
-    return computed_md5_hash == original_md5_hash
+    return computed_md5_hash.hexdigest() == original_md5_hash
 
 
 def download_assets(asset_name: Union[VideoAssets, str]) -> str:

supervision/config.py

@@ -0,0 +1 @@
+CLASS_NAME_DATA_FIELD = "class_name"

supervision/dataset/formats/pascal_voc.py

@@ -268,12 +268,12 @@ def detections_from_xml_obj(
         [extended_classes.index(class_name) for class_name in class_names]
     )
 
-    if with_masks:
-        annotation = Detections(
-            xyxy=xyxy, mask=np.array(masks).astype(bool), class_id=class_id
-        )
-    else:
-        annotation = Detections(xyxy=xyxy, class_id=class_id)
+    annotation = Detections(
+        xyxy=xyxy.astype(np.float32),
+        mask=np.array(masks).astype(bool) if with_masks else None,
+        class_id=class_id,
+    )
+
     return annotation, extended_classes
 
 

supervision/detection/core.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from contextlib import suppress
-from dataclasses import astuple, dataclass
+from dataclasses import dataclass, field
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
@@ -9,63 +9,22 @@ import numpy as np
 from supervision.detection.utils import (
     calculate_masks_centroids,
     extract_ultralytics_masks,
+    get_data_item,
+    is_data_equal,
+    merge_data,
     non_max_suppression,
     process_roboflow_result,
+    validate_detections_fields,
     xywh_to_xyxy,
 )
 from supervision.geometry.core import Position
 from supervision.utils.internal import deprecated
 
 
-def _validate_xyxy(xyxy: Any, n: int) -> None:
-    is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
-    if not is_valid:
-        raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
-
-
-def _validate_mask(mask: Any, n: int) -> None:
-    is_valid = mask is None or (
-        isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
-    )
-    if not is_valid:
-        raise ValueError("mask must be 3d np.ndarray with (n, H, W) shape")
-
-
-def validate_inference_callback(callback) -> None:
-    tmp_img = np.zeros((256, 256, 3), dtype=np.uint8)
-    res = callback(tmp_img)
-    if not isinstance(res, Detections):
-        raise ValueError("Callback function must return sv.Detection type")
-
-
-def _validate_class_id(class_id: Any, n: int) -> None:
-    is_valid = class_id is None or (
-        isinstance(class_id, np.ndarray) and class_id.shape == (n,)
-    )
-    if not is_valid:
-        raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
-
-
-def _validate_confidence(confidence: Any, n: int) -> None:
-    is_valid = confidence is None or (
-        isinstance(confidence, np.ndarray) and confidence.shape == (n,)
-    )
-    if not is_valid:
-        raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
-
-
-def _validate_tracker_id(tracker_id: Any, n: int) -> None:
-    is_valid = tracker_id is None or (
-        isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
-    )
-    if not is_valid:
-        raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
-
-
 @dataclass
 class Detections:
     """
-    Data class containing information about the detections in a video frame.
+    A dataclass representing detection results.
 
     Attributes:
         xyxy (np.ndarray): An array of shape `(n, 4)` containing
@@ -78,6 +37,17 @@ class Detections:
             `(n,)` containing the class ids of the detections.
         tracker_id (Optional[np.ndarray]): An array of shape
             `(n,)` containing the tracker ids of the detections.
+        data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional
+            data where each key is a string representing the data type, and the value
+            is either a NumPy array or a list of corresponding data.
+
+    !!! warning
+
+        The `data` field in the `sv.Detections` class is currently in an experimental
+        phase. Please be aware that its API and functionality are subject to change in
+        future updates as we continue to refine and improve its capabilities.
+        We encourage users to experiment with this feature and provide feedback, but
+        also to be prepared for potential modifications in upcoming releases.
     """
 
     xyxy: np.ndarray
@@ -85,14 +55,17 @@ class Detections:
     confidence: Optional[np.ndarray] = None
     class_id: Optional[np.ndarray] = None
     tracker_id: Optional[np.ndarray] = None
+    data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
 
     def __post_init__(self):
-        n = len(self.xyxy)
-        _validate_xyxy(xyxy=self.xyxy, n=n)
-        _validate_mask(mask=self.mask, n=n)
-        _validate_class_id(class_id=self.class_id, n=n)
-        _validate_confidence(confidence=self.confidence, n=n)
-        _validate_tracker_id(tracker_id=self.tracker_id, n=n)
+        validate_detections_fields(
+            xyxy=self.xyxy,
+            mask=self.mask,
+            confidence=self.confidence,
+            class_id=self.class_id,
+            tracker_id=self.tracker_id,
+            data=self.data,
+        )
 
     def __len__(self):
         """
@@ -109,11 +82,12 @@ class Detections:
             Optional[float],
             Optional[int],
             Optional[int],
+            Dict[str, Union[np.ndarray, List]],
         ]
     ]:
         """
         Iterates over the Detections object and yield a tuple of
-        `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
+        `(xyxy, mask, confidence, class_id, tracker_id, data)` for each detection.
         """
         for i in range(len(self.xyxy)):
             yield (
@@ -122,36 +96,18 @@ class Detections:
                 self.confidence[i] if self.confidence is not None else None,
                 self.class_id[i] if self.class_id is not None else None,
                 self.tracker_id[i] if self.tracker_id is not None else None,
+                get_data_item(self.data, i),
             )
 
     def __eq__(self, other: Detections):
         return all(
             [
                 np.array_equal(self.xyxy, other.xyxy),
-                any(
-                    [
-                        self.mask is None and other.mask is None,
-                        np.array_equal(self.mask, other.mask),
-                    ]
-                ),
-                any(
-                    [
-                        self.class_id is None and other.class_id is None,
-                        np.array_equal(self.class_id, other.class_id),
-                    ]
-                ),
-                any(
-                    [
-                        self.confidence is None and other.confidence is None,
-                        np.array_equal(self.confidence, other.confidence),
-                    ]
-                ),
-                any(
-                    [
-                        self.tracker_id is None and other.tracker_id is None,
-                        np.array_equal(self.tracker_id, other.tracker_id),
-                    ]
-                ),
+                np.array_equal(self.mask, other.mask),
+                np.array_equal(self.class_id, other.class_id),
+                np.array_equal(self.confidence, other.confidence),
+                np.array_equal(self.tracker_id, other.tracker_id),
+                is_data_equal(self.data, other.data),
             ]
         )
 
@@ -443,11 +399,19 @@ class Detections:
     def from_inference(cls, roboflow_result: Union[dict, Any]) -> Detections:
         """
         Create a Detections object from the [Roboflow](https://roboflow.com/)
-            API inference result or the [Inference](https://inference.roboflow.com/) package results.
+        API inference result or the [Inference](https://inference.roboflow.com/)
+        package results. This method extracts bounding boxes, class IDs,
+        confidences, and class names from the Roboflow API result and encapsulates
+        them into a Detections object.
+
+        !!! note
+
+            Class names can be accessed using the key 'class_name' in the returned
+            object's data attribute.
 
         Args:
             roboflow_result (dict, any): The result from the
-                Roboflow API or Inference package containing predictions
+                Roboflow API or Inference package containing predictions.
 
         Returns:
             (Detections): A Detections object containing the bounding boxes, class IDs,
@@ -473,11 +437,12 @@ class Detections:
             ... }
 
             >>> detections = sv.Detections.from_roboflow(roboflow_result)
+            >>> detections['class_name']
             ```
         """
         with suppress(AttributeError):
             roboflow_result = roboflow_result.dict(exclude_none=True, by_alias=True)
-        xyxy, confidence, class_id, masks, trackers = process_roboflow_result(
+        xyxy, confidence, class_id, masks, trackers, data = process_roboflow_result(
             roboflow_result=roboflow_result
         )
 
@@ -490,6 +455,7 @@ class Detections:
             class_id=class_id,
             mask=masks,
             tracker_id=trackers,
+            data=data,
         )
 
     @classmethod
@@ -744,30 +710,67 @@ class Detections:
 
         Example:
             ```python
-            >>> from supervision import Detections
+            import numpy as np
+            import supervision as sv
 
-            >>> detections_1 = Detections(...)
-            >>> detections_2 = Detections(...)
+            >>> detections_1 = sv.Detections(
+            ...     xyxy=np.array([[15, 15, 100, 100], [200, 200, 300, 300]]),
+            ...     class_id=np.array([1, 2]),
+            ...     data={'feature_vector': np.array([0.1, 0.2)])}
+            ... )
+
+            >>> detections_2 = sv.Detections(
+            ...     xyxy=np.array([[30, 30, 120, 120]]),
+            ...     class_id=np.array([1]),
+            ...     data={'feature_vector': [np.array([0.3])]}
+            ... )
 
             >>> merged_detections = Detections.merge([detections_1, detections_2])
+
+            >>> merged_detections.xyxy
+            array([[ 15,  15, 100, 100],
+                   [200, 200, 300, 300],
+                   [ 30,  30, 120, 120]])
+
+            >>> merged_detections.class_id
+            array([1, 2, 1])
+
+            >>> merged_detections.data['feature_vector']
+            array([0.1, 0.2, 0.3])
             ```
         """
         if len(detections_list) == 0:
             return Detections.empty()
 
-        detections_tuples_list = [astuple(detection) for detection in detections_list]
-        xyxy, mask, confidence, class_id, tracker_id = [
-            list(field) for field in zip(*detections_tuples_list)
-        ]
+        for detections in detections_list:
+            validate_detections_fields(
+                xyxy=detections.xyxy,
+                mask=detections.mask,
+                confidence=detections.confidence,
+                class_id=detections.class_id,
+                tracker_id=detections.tracker_id,
+                data=detections.data,
+            )
 
-        def __all_not_none(item_list: List[Any]):
-            return all(x is not None for x in item_list)
+        xyxy = np.vstack([d.xyxy for d in detections_list])
 
-        xyxy = np.vstack(xyxy)
-        mask = np.vstack(mask) if __all_not_none(mask) else None
-        confidence = np.hstack(confidence) if __all_not_none(confidence) else None
-        class_id = np.hstack(class_id) if __all_not_none(class_id) else None
-        tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None
+        def stack_or_none(name: str):
+            if all(d.__getattribute__(name) is None for d in detections_list):
+                return None
+            if any(d.__getattribute__(name) is None for d in detections_list):
+                raise ValueError(f"All or none of the '{name}' fields must be None")
+            return (
+                np.vstack([d.__getattribute__(name) for d in detections_list])
+                if name == "mask"
+                else np.hstack([d.__getattribute__(name) for d in detections_list])
+            )
+
+        mask = stack_or_none("mask")
+        confidence = stack_or_none("confidence")
+        class_id = stack_or_none("class_id")
+        tracker_id = stack_or_none("tracker_id")
+
+        data = merge_data([d.data for d in detections_list])
 
         return cls(
             xyxy=xyxy,
@@ -775,6 +778,7 @@ class Detections:
             confidence=confidence,
             class_id=class_id,
             tracker_id=tracker_id,
+            data=data,
         )
 
     def get_anchors_coordinates(self, anchor: Position) -> np.ndarray:
@@ -844,17 +848,23 @@ class Detections:
         raise ValueError(f"{anchor} is not supported.")
 
     def __getitem__(
-        self, index: Union[int, slice, List[int], np.ndarray]
-    ) -> Detections:
+        self, index: Union[int, slice, List[int], np.ndarray, str]
+    ) -> Union[Detections, List, np.ndarray, None]:
         """
-        Get a subset of the Detections object.
+        Get a subset of the Detections object or access an item from its data field.
+
+        When provided with an integer, slice, list of integers, or a numpy array, this
+        method returns a new Detections object that represents a subset of the original
+        detections. When provided with a string, it accesses the corresponding item in
+        the data dictionary.
 
         Args:
-            index (Union[int, slice, List[int], np.ndarray]):
-                The index or indices of the subset of the Detections
+            index (Union[int, slice, List[int], np.ndarray, str]): The index, indices,
+                or key to access a subset of the Detections or an item from the data.
 
         Returns:
-            (Detections): A subset of the Detections object.
+            Union[Detections, Any]: A subset of the Detections object or an item from
+                the data field.
 
         Example:
             ```python
@@ -863,16 +873,16 @@ class Detections:
             >>> detections = sv.Detections(...)
 
             >>> first_detection = detections[0]
-
             >>> first_10_detections = detections[0:10]
-
             >>> some_detections = detections[[0, 2, 4]]
-
             >>> class_0_detections = detections[detections.class_id == 0]
-
             >>> high_confidence_detections = detections[detections.confidence > 0.5]
+
+            >>> feature_vector = detections['feature_vector']
             ```
         """
+        if isinstance(index, str):
+            return self.data.get(index)
         if isinstance(index, int):
             index = [index]
         return Detections(
@@ -881,8 +891,45 @@ class Detections:
             confidence=self.confidence[index] if self.confidence is not None else None,
             class_id=self.class_id[index] if self.class_id is not None else None,
             tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
+            data=get_data_item(self.data, index),
         )
 
+    def __setitem__(self, key: str, value: Union[np.ndarray, List]):
+        """
+        Set a value in the data dictionary of the Detections object.
+
+        Args:
+            key (str): The key in the data dictionary to set.
+            value (Union[np.ndarray, List]): The value to set for the key.
+
+        Example:
+            ```python
+            >>> import cv2
+            >>> from ultralytics import YOLO
+            >>> import supervision as sv
+
+            >>> model = YOLO('yolov8s.pt')
+
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
+
+            >>> result = model(image)[0]
+            >>> detections = sv.Detections.from_ultralytics(result)
+
+            >>> detections['names'] = [
+            ...     model.model.names[class_id]
+            ...     for class_id
+            ...     in detections.class_id
+            ... ]
+            ```
+        """
+        if not isinstance(value, (np.ndarray, list)):
+            raise TypeError("Value must be a np.ndarray or a list")
+
+        if isinstance(value, list):
+            value = np.array(value)
+
+        self.data[key] = value
+
     @property
     def area(self) -> np.ndarray:
         """

supervision/detection/tools/inference_slicer.py

@@ -3,7 +3,7 @@ from typing import Callable, Optional, Tuple
 
 import numpy as np
 
-from supervision.detection.core import Detections, validate_inference_callback
+from supervision.detection.core import Detections
 from supervision.detection.utils import move_boxes
 from supervision.utils.image import crop_image
 
@@ -60,7 +60,6 @@ class InferenceSlicer:
         self.iou_threshold = iou_threshold
         self.callback = callback
         self.thread_workers = thread_workers
-        validate_inference_callback(callback=callback)
 
     def __call__(self, image: np.ndarray) -> Detections:
         """

supervision/detection/utils.py

@@ -1,8 +1,11 @@
-from typing import List, Optional, Tuple
+from itertools import chain
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
 
+from supervision.config import CLASS_NAME_DATA_FIELD
+
 MIN_POLYGON_POINT_COUNT = 3
 
 
@@ -299,7 +302,6 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
     orig_shape = yolov8_results.orig_shape
     inference_shape = tuple(yolov8_results.masks.data.shape[1:])
 
-    gain = 0
     pad = (0, 0)
 
     if inference_shape != orig_shape:
@@ -331,13 +333,28 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
 
 def process_roboflow_result(
     roboflow_result: dict,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
+) -> Tuple[
+    np.ndarray,
+    np.ndarray,
+    np.ndarray,
+    Optional[np.ndarray],
+    Optional[np.ndarray],
+    Dict[str, List[np.ndarray]],
+]:
     if not roboflow_result["predictions"]:
-        return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
+        return (
+            np.empty((0, 4)),
+            np.empty(0),
+            np.empty(0),
+            None,
+            None,
+            {CLASS_NAME_DATA_FIELD: np.empty(0)},
+        )
 
     xyxy = []
     confidence = []
     class_id = []
+    class_name = []
     masks = []
     tracker_ids = []
 
@@ -357,6 +374,7 @@ def process_roboflow_result(
         if "points" not in prediction:
             xyxy.append([x_min, y_min, x_max, y_max])
             class_id.append(prediction["class_id"])
+            class_name.append(prediction["class"])
             confidence.append(prediction["confidence"])
             if "tracker_id" in prediction:
                 tracker_ids.append(prediction["tracker_id"])
@@ -367,6 +385,7 @@ def process_roboflow_result(
             mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
             xyxy.append([x_min, y_min, x_max, y_max])
             class_id.append(prediction["class_id"])
+            class_name.append(prediction["class"])
             confidence.append(prediction["confidence"])
             masks.append(mask)
             if "tracker_id" in prediction:
@@ -375,10 +394,12 @@ def process_roboflow_result(
     xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
     confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0)
     class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0)
+    class_name = np.array(class_name) if len(class_name) > 0 else np.empty(0)
     masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
     tracker_id = np.array(tracker_ids).astype(int) if len(tracker_ids) > 0 else None
+    data = {CLASS_NAME_DATA_FIELD: class_name}
 
-    return xyxy, confidence, class_id, masks, tracker_id
+    return xyxy, confidence, class_id, masks, tracker_id, data
 
 
 def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
@@ -470,3 +491,204 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray:
     centroid_y = sum_over_mask(vertical_indices, aggregation_axis) / total_pixels
 
     return np.column_stack((centroid_x, centroid_y)).astype(int)
+
+
+def validate_xyxy(xyxy: Any) -> None:
+    expected_shape = "(_, 4)"
+    actual_shape = str(getattr(xyxy, "shape", None))
+    is_valid = isinstance(xyxy, np.ndarray) and xyxy.ndim == 2 and xyxy.shape[1] == 4
+    if not is_valid:
+        raise ValueError(
+            f"xyxy must be a 2D np.ndarray with shape {expected_shape}, but got shape "
+            f"{actual_shape}"
+        )
+
+
+def validate_mask(mask: Any, n: int) -> None:
+    expected_shape = f"({n}, H, W)"
+    actual_shape = str(getattr(mask, "shape", None))
+    is_valid = mask is None or (
+        isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
+    )
+    if not is_valid:
+        raise ValueError(
+            f"mask must be a 3D np.ndarray with shape {expected_shape}, but got shape "
+            f"{actual_shape}"
+        )
+
+
+def validate_class_id(class_id: Any, n: int) -> None:
+    expected_shape = f"({n},)"
+    actual_shape = str(getattr(class_id, "shape", None))
+    is_valid = class_id is None or (
+        isinstance(class_id, np.ndarray) and class_id.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError(
+            f"class_id must be a 1D np.ndarray with shape {expected_shape}, but got "
+            f"shape {actual_shape}"
+        )
+
+
+def validate_confidence(confidence: Any, n: int) -> None:
+    expected_shape = f"({n},)"
+    actual_shape = str(getattr(confidence, "shape", None))
+    is_valid = confidence is None or (
+        isinstance(confidence, np.ndarray) and confidence.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError(
+            f"confidence must be a 1D np.ndarray with shape {expected_shape}, but got "
+            f"shape {actual_shape}"
+        )
+
+
+def validate_tracker_id(tracker_id: Any, n: int) -> None:
+    expected_shape = f"({n},)"
+    actual_shape = str(getattr(tracker_id, "shape", None))
+    is_valid = tracker_id is None or (
+        isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError(
+            f"tracker_id must be a 1D np.ndarray with shape {expected_shape}, but got "
+            f"shape {actual_shape}"
+        )
+
+
+def validate_data(data: Dict[str, Any], n: int) -> None:
+    for key, value in data.items():
+        if isinstance(value, list):
+            if len(value) != n:
+                raise ValueError(f"Length of list for key '{key}' must be {n}")
+        elif isinstance(value, np.ndarray):
+            if value.ndim == 1 and value.shape[0] != n:
+                raise ValueError(f"Shape of np.ndarray for key '{key}' must be ({n},)")
+            elif value.ndim > 1 and value.shape[0] != n:
+                raise ValueError(
+                    f"First dimension of np.ndarray for key '{key}' must have size {n}"
+                )
+        else:
+            raise ValueError(f"Value for key '{key}' must be a list or np.ndarray")
+
+
+def validate_detections_fields(
+    xyxy: Any,
+    mask: Any,
+    class_id: Any,
+    confidence: Any,
+    tracker_id: Any,
+    data: Dict[str, Any],
+) -> None:
+    validate_xyxy(xyxy)
+    n = len(xyxy)
+    validate_mask(mask, n)
+    validate_class_id(class_id, n)
+    validate_confidence(confidence, n)
+    validate_tracker_id(tracker_id, n)
+    validate_data(data, n)
+
+
+def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray]) -> bool:
+    """
+    Compares the data payloads of two Detections instances.
+
+    Args:
+        data_a, data_b: The data payloads of the instances.
+
+    Returns:
+        True if the data payloads are equal, False otherwise.
+    """
+    return set(data_a.keys()) == set(data_b.keys()) and all(
+        np.array_equal(data_a[key], data_b[key]) for key in data_a
+    )
+
+
+def merge_data(
+    data_list: List[Dict[str, Union[np.ndarray, List]]],
+) -> Dict[str, Union[np.ndarray, List]]:
+    """
+    Merges the data payloads of a list of Detections instances.
+
+    Args:
+        data_list: The data payloads of the instances.
+
+    Returns:
+        A single data payload containing the merged data, preserving the original data
+            types (list or np.ndarray).
+
+    Raises:
+        ValueError: If data values within a single object have different lengths or if
+            dictionaries have different keys.
+    """
+    if not data_list:
+        return {}
+
+    all_keys_sets = [set(data.keys()) for data in data_list]
+    if not all(keys_set == all_keys_sets[0] for keys_set in all_keys_sets):
+        raise ValueError("All data dictionaries must have the same keys to merge.")
+
+    for data in data_list:
+        lengths = [len(value) for value in data.values()]
+        if len(set(lengths)) > 1:
+            raise ValueError(
+                "All data values within a single object must have equal length."
+            )
+
+    merged_data = {key: [] for key in all_keys_sets[0]}
+
+    for data in data_list:
+        for key in merged_data:
+            merged_data[key].append(data[key])
+
+    for key in merged_data:
+        if all(isinstance(item, list) for item in merged_data[key]):
+            merged_data[key] = list(chain.from_iterable(merged_data[key]))
+        elif all(isinstance(item, np.ndarray) for item in merged_data[key]):
+            ndim = merged_data[key][0].ndim
+            if ndim == 1:
+                merged_data[key] = np.hstack(merged_data[key])
+            elif ndim > 1:
+                merged_data[key] = np.vstack(merged_data[key])
+            else:
+                raise ValueError(f"Unexpected array dimension for key '{key}'.")
+        else:
+            raise ValueError(
+                f"Inconsistent data types for key '{key}'. Only np.ndarray and list "
+                f"types are allowed."
+            )
+
+    return merged_data
+
+
+def get_data_item(
+    data: Dict[str, Union[np.ndarray, List]],
+    index: Union[int, slice, List[int], np.ndarray],
+) -> Dict[str, Union[np.ndarray, List]]:
+    """
+    Retrieve a subset of the data dictionary based on the given index.
+
+    Args:
+        data: The data dictionary of the Detections object.
+        index: The index or indices specifying the subset to retrieve.
+
+    Returns:
+        A subset of the data dictionary corresponding to the specified index.
+    """
+    subset_data = {}
+    for key, value in data.items():
+        if isinstance(value, np.ndarray):
+            subset_data[key] = value[index]
+        elif isinstance(value, list):
+            if isinstance(index, slice):
+                subset_data[key] = value[index]
+            elif isinstance(index, (list, np.ndarray)):
+                subset_data[key] = [value[i] for i in index]
+            elif isinstance(index, int):
+                subset_data[key] = [value[index]]
+            else:
+                raise TypeError(f"Unsupported index type: {type(index)}")
+        else:
+            raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
+
+    return subset_data

test/dataset/formats/test_pascal_voc.py

@@ -122,7 +122,7 @@ NO_DETECTIONS = """<annotation></annotation>"""
             ["test"],
             (100, 100),
             False,
-            mock_detections(np.array([[0, 0, 10, 10]]), None, [0]),
+            mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
             DoesNotRaise(),
         ),
         (
@@ -130,7 +130,9 @@ NO_DETECTIONS = """<annotation></annotation>"""
             ["test"],
             (100, 100),
             False,
-            mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]),
+            mock_detections(
+                xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), class_id=[0, 0]
+            ),
             DoesNotRaise(),
         ),
         (
@@ -139,9 +141,8 @@ NO_DETECTIONS = """<annotation></annotation>"""
             (100, 100),
             False,
             mock_detections(
-                np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
-                None,
-                [0, 0, 1],
+                xyxy=np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
+                class_id=[0, 0, 1],
             ),
             DoesNotRaise(),
         ),
@@ -150,7 +151,7 @@ NO_DETECTIONS = """<annotation></annotation>"""
             [],
             (100, 100),
             False,
-            mock_detections(np.empty((0, 4)), None, []),
+            mock_detections(xyxy=np.empty((0, 4)), class_id=[]),
             DoesNotRaise(),
         ),
     ],

test/detection/test_core.py

@@ -153,7 +153,10 @@ def test_getitem(
             DoesNotRaise(),
         ),  # single detection with xyxy field
         (
-            [mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
+            [
+                mock_detections(xyxy=[[10, 10, 20, 20]]),
+                mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)),
+            ],
             mock_detections(xyxy=[[10, 10, 20, 20]]),
             DoesNotRaise(),
         ),  # single detection with xyxy field + empty detection
@@ -171,7 +174,7 @@ def test_getitem(
                 mock_detections(xyxy=[[20, 20, 30, 30]]),
             ],
             mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
-            DoesNotRaise(),
+            pytest.raises(ValueError),
         ),  # detection with xyxy, class_id fields + detection with xyxy field
         (
             [
@@ -181,6 +184,16 @@ def test_getitem(
             mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]),
             DoesNotRaise(),
         ),  # two detections with xyxy, class_id fields
+        (
+            [
+                mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+                mock_detections(xyxy=[[20, 20, 30, 30]], data={"test": [2]}),
+            ],
+            mock_detections(
+                xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], data={"test": [1, 2]}
+            ),
+            DoesNotRaise(),
+        ),  # two detections with xyxy, data fields
     ],
 )
 def test_merge(
@@ -273,3 +286,54 @@ def test_get_anchor_coordinates(
     result = detections.get_anchors_coordinates(anchor)
     with exception:
         assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    "detections_a, detections_b, expected_result",
+    [
+        (
+            Detections.empty(),
+            Detections.empty(),
+            True,
+        ),  # empty detections
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            True,
+        ),  # detections with xyxy field
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+            mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+            True,
+        ),  # detections with xyxy, confidence fields
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            False,
+        ),  # detection with xyxy field + detection with xyxy, confidence fields
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+            True,
+        ),  # detections with xyxy, data fields
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            False,
+        ),  # detection with xyxy field + detection with xyxy, data fields
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_2": [1]}),
+            False,
+        ),  # detections with xyxy, and different data field names
+        (
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
+            mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [3]}),
+            False,
+        ),  # detections with xyxy, and different data field values
+    ],
+)
+def test_equal(
+    detections_a: Detections, detections_b: Detections, expected_result: bool
+) -> None:
+    assert (detections_a == detections_b) == expected_result

test/detection/test_utils.py

@@ -1,13 +1,16 @@
 from contextlib import ExitStack as DoesNotRaise
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import pytest
 
+from supervision.config import CLASS_NAME_DATA_FIELD
 from supervision.detection.utils import (
     calculate_masks_centroids,
     clip_boxes,
     filter_polygons_by_area,
+    get_data_item,
+    merge_data,
     move_boxes,
     non_max_suppression,
     process_roboflow_result,
@@ -266,7 +269,14 @@ def test_filter_polygons_by_area(
     [
         (
             {"predictions": [], "image": {"width": 1000, "height": 1000}},
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {CLASS_NAME_DATA_FIELD: np.empty(0)},
+            ),
             DoesNotRaise(),
         ),  # empty result
         (
@@ -290,6 +300,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 None,
                 None,
+                {CLASS_NAME_DATA_FIELD: np.array(["person"])},
             ),
             DoesNotRaise(),
         ),  # single correct object detection result
@@ -325,6 +336,7 @@ def test_filter_polygons_by_area(
                 np.array([0, 7]),
                 None,
                 np.array([1, 2]),
+                {CLASS_NAME_DATA_FIELD: np.array(["person", "truck"])},
             ),
             DoesNotRaise(),
         ),  # two correct object detection result
@@ -345,7 +357,14 @@ def test_filter_polygons_by_area(
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {CLASS_NAME_DATA_FIELD: np.empty(0)},
+            ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no points
         (
@@ -364,7 +383,14 @@ def test_filter_polygons_by_area(
                 ],
                 "image": {"width": 1000, "height": 1000},
             },
-            (np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None,
+                None,
+                {CLASS_NAME_DATA_FIELD: np.empty(0)},
+            ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
         (
@@ -394,6 +420,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 TEST_MASK,
                 None,
+                {CLASS_NAME_DATA_FIELD: np.array(["person"])},
             ),
             DoesNotRaise(),
         ),  # single incorrect instance segmentation result with no enough points
@@ -434,6 +461,7 @@ def test_filter_polygons_by_area(
                 np.array([0]),
                 TEST_MASK,
                 None,
+                {CLASS_NAME_DATA_FIELD: np.array(["person"])},
             ),
             DoesNotRaise(),
         ),  # two instance segmentation results - one correct, one incorrect
@@ -457,6 +485,15 @@ def test_process_roboflow_result(
         assert (result[4] is None and expected_result[4] is None) or (
             np.array_equal(result[4], expected_result[4])
         )
+        for key in result[5]:
+            if isinstance(result[5][key], np.ndarray):
+                assert np.array_equal(
+                    result[5][key], expected_result[5][key]
+                ), f"Mismatch in arrays for key {key}"
+            else:
+                assert (
+                    result[5][key] == expected_result[5][key]
+                ), f"Mismatch in non-array data for key {key}"
 
 
 @pytest.mark.parametrize(
@@ -640,3 +677,274 @@ def test_calculate_masks_centroids(
     with exception:
         result = calculate_masks_centroids(masks=masks)
         assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    "data_list, expected_result, exception",
+    [
+        (
+            [],
+            {},
+            DoesNotRaise(),
+        ),  # empty data list
+        (
+            [{}],
+            {},
+            DoesNotRaise(),
+        ),  # single empty data dict
+        (
+            [{}, {}],
+            {},
+            DoesNotRaise(),
+        ),  # two empty data dicts
+        (
+            [
+                {"test_1": []},
+            ],
+            {"test_1": []},
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and empty list values
+        (
+            [
+                {"test_1": np.array([])},
+            ],
+            {"test_1": np.array([])},
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and empty np.array values
+        (
+            [
+                {"test_1": [1, 2, 3]},
+            ],
+            {"test_1": [1, 2, 3]},
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            [
+                {"test_1": []},
+                {"test_1": [3, 2, 1]},
+            ],
+            {"test_1": [3, 2, 1]},
+            DoesNotRaise(),
+        ),  # two data dicts with the same field name and empty and list values
+        (
+            [
+                {"test_1": [1, 2, 3]},
+                {"test_1": [3, 2, 1]},
+            ],
+            {"test_1": [1, 2, 3, 3, 2, 1]},
+            DoesNotRaise(),
+        ),  # two data dicts with the same field name and list values
+        (
+            [
+                {"test_1": [1, 2, 3]},
+                {"test_1": [3, 2, 1]},
+                {"test_1": [1, 2, 3]},
+            ],
+            {"test_1": [1, 2, 3, 3, 2, 1, 1, 2, 3]},
+            DoesNotRaise(),
+        ),  # three data dicts with the same field name and list values
+        (
+            [
+                {"test_1": [1, 2, 3]},
+                {"test_2": [3, 2, 1]},
+            ],
+            None,
+            pytest.raises(ValueError),
+        ),  # two data dicts with different field names
+        (
+            [
+                {"test_1": np.array([1, 2, 3])},
+                {"test_1": np.array([3, 2, 1])},
+            ],
+            {"test_1": np.array([1, 2, 3, 3, 2, 1])},
+            DoesNotRaise(),
+        ),  # two data dicts with the same field name and np.array values as 1D arrays
+        (
+            [
+                {"test_1": np.array([[1, 2, 3]])},
+                {"test_1": np.array([[3, 2, 1]])},
+            ],
+            {"test_1": np.array([[1, 2, 3], [3, 2, 1]])},
+            DoesNotRaise(),
+        ),  # two data dicts with the same field name and np.array values as 2D arrays
+        (
+            [
+                {"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b", "c"])},
+                {"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
+            ],
+            {
+                "test_1": np.array([1, 2, 3, 3, 2, 1]),
+                "test_2": np.array(["a", "b", "c", "c", "b", "a"]),
+            },
+            DoesNotRaise(),
+        ),  # two data dicts with the same field names and np.array values
+        (
+            [
+                {"test_1": [1, 2, 3], "test_2": np.array(["a", "b", "c"])},
+                {"test_1": [3, 2, 1], "test_2": np.array(["c", "b", "a"])},
+            ],
+            {
+                "test_1": [1, 2, 3, 3, 2, 1],
+                "test_2": np.array(["a", "b", "c", "c", "b", "a"]),
+            },
+            DoesNotRaise(),
+        ),  # two data dicts with the same field names and mixed values
+        (
+            [
+                {"test_1": np.array([1, 2, 3])},
+                {"test_1": np.array([[3, 2, 1]])},
+            ],
+            None,
+            pytest.raises(ValueError),
+        ),  # two data dicts with the same field name and 1D and 2D arrays values
+        (
+            [
+                {"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b"])},
+                {"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
+            ],
+            None,
+            pytest.raises(ValueError),
+        ),  # two data dicts with the same field name and different length arrays values
+    ],
+)
+def test_merge_data(
+    data_list: List[Dict[str, Any]],
+    expected_result: Optional[Dict[str, Any]],
+    exception: Exception,
+):
+    with exception:
+        result = merge_data(data_list=data_list)
+        for key in result:
+            if isinstance(result[key], np.ndarray):
+                assert np.array_equal(
+                    result[key], expected_result[key]
+                ), f"Mismatch in arrays for key {key}"
+            else:
+                assert (
+                    result[key] == expected_result[key]
+                ), f"Mismatch in non-array data for key {key}"
+
+
+@pytest.mark.parametrize(
+    "data, index, expected_result, exception",
+    [
+        ({}, 0, {}, DoesNotRaise()),  # empty data dict
+        (
+            {
+                "test_1": [1, 2, 3],
+            },
+            0,
+            {
+                "test_1": [1],
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            {
+                "test_1": np.array([1, 2, 3]),
+            },
+            0,
+            {
+                "test_1": np.array([1]),
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and np.array values as 1D arrays
+        (
+            {
+                "test_1": [1, 2, 3],
+            },
+            slice(0, 2),
+            {
+                "test_1": [1, 2],
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            {
+                "test_1": np.array([1, 2, 3]),
+            },
+            slice(0, 2),
+            {
+                "test_1": np.array([1, 2]),
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and np.array values as 1D arrays
+        (
+            {
+                "test_1": [1, 2, 3],
+            },
+            -1,
+            {
+                "test_1": [3],
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            {
+                "test_1": np.array([1, 2, 3]),
+            },
+            -1,
+            {
+                "test_1": np.array([3]),
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and np.array values as 1D arrays
+        (
+            {
+                "test_1": [1, 2, 3],
+            },
+            [0, 2],
+            {
+                "test_1": [1, 3],
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            {
+                "test_1": np.array([1, 2, 3]),
+            },
+            [0, 2],
+            {
+                "test_1": np.array([1, 3]),
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and np.array values as 1D arrays
+        (
+            {
+                "test_1": [1, 2, 3],
+            },
+            np.array([0, 2]),
+            {
+                "test_1": [1, 3],
+            },
+            DoesNotRaise(),
+        ),  # single data dict with a single field name and list values
+        (
+            {
+                "test_1": np.array([1, 2, 3]),
+            },
+            np.array([0, 2]),
+            {
+                "test_1": np.array([1, 3]),
+            },
+            DoesNotRaise(),
+        ),
+    ],
+)
+def test_get_data_item(
+    data: Dict[str, Any],
+    index: Any,
+    expected_result: Optional[Dict[str, Any]],
+    exception: Exception,
+):
+    with exception:
+        result = get_data_item(data=data, index=index)
+        for key in result:
+            if isinstance(result[key], np.ndarray):
+                assert np.array_equal(
+                    result[key], expected_result[key]
+                ), f"Mismatch in arrays for key {key}"
+            else:
+                assert (
+                    result[key] == expected_result[key]
+                ), f"Mismatch in non-array data for key {key}"

test/test_utils.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 
@@ -7,19 +7,26 @@ from supervision.detection.core import Detections
 
 def mock_detections(
     xyxy: List[List[float]],
-    confidence: List[float] = None,
-    class_id: List[int] = None,
-    tracker_id: List[int] = None,
+    mask: Optional[List[np.ndarray]] = None,
+    confidence: Optional[List[float]] = None,
+    class_id: Optional[List[int]] = None,
+    tracker_id: Optional[List[int]] = None,
+    data: Optional[Dict[str, List[Any]]] = None,
 ) -> Detections:
+    def convert_data(data: Dict[str, List[Any]]):
+        return {k: np.array(v) for k, v in data.items()}
+
     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),
+        mask=(mask if mask is None else np.array(mask, dtype=bool)),
+        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)
+        ),
+        data=convert_data(data) if data else {},
     )