Add AWS Rekognition CustomLabels data loader

supervision/detection/core.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import astuple, dataclass
-from typing import Any, Iterator, List, Optional, Tuple, Union
+from typing import Any, Iterator, List, Optional, Tuple, Union, Dict
 
 import numpy as np
 
@@ -483,6 +483,89 @@ def from_sam(cls, sam_result: List[dict]) -> Detections:
         xyxy = xywh_to_xyxy(boxes_xywh=xywh)
         return cls(xyxy=xyxy, mask=mask)
 
+    @classmethod
+    def from_rekognition_detectlabels(cls, rekognition_det: dict, imgsz: Tuple[int, int], class_map: Optional[Dict[int, str]] = None) -> Detections:
+        """
+        Creates a Detections instance from
+            AWS Rekognition DetectLabels
+            (https://docs.aws.amazon.com/rekognition/latest/dg/labels-detect-labels-image.html)
+            inference result.
+
+        Args:
+            rekognition_det (List[dict]): The output from AWS Rekognition DetectLabels
+
+        Returns:
+            Detections: A new Detections object.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+            >>> import boto3
+            >>> import io
+            >>> from PIL import Image
+
+            >>> session = boto3.Session()
+            >>> client = session.client("rekognition")
+
+            >>> with Image.open(input) as image:
+            >>>    buffered = io.BytesIO()
+            >>>    image.save(buffered, format=image.format)
+            >>>    image_bytes = buffered.getvalue()
+
+            >>> response = client.detect_labels(Image={"Bytes": image_bytes})
+
+            >>> detections = sv.Detections.from_rekognition_detectlabels(response)
+            ```
+        """
+
+        xyxys, confidences, class_ids = [], [], []
+
+        is_dynamic_mapping = class_map is None
+
+        if is_dynamic_mapping:
+            class_map = {}
+
+        class_map = {value: key for key, value in class_map.items()}
+
+        for label in rekognition_det["Labels"]:
+            if len(label["Instances"]) == 0:
+                continue
+
+            for box in label["Instances"]:
+                x0 = box["BoundingBox"]["Left"]
+                y0 = box["BoundingBox"]["Top"]
+                x1 = x0 + box["BoundingBox"]["Width"]
+                y1 = y0 + box["BoundingBox"]["Height"]
+
+                x0 *= imgsz[0]
+                y0 *= imgsz[1]
+
+                x1 *= imgsz[0]
+                y1 *= imgsz[1]
+
+                class_name = label["Name"]
+                confidence = box["Confidence"]
+
+                class_id = class_map.get(class_name, None)
+
+                if is_dynamic_mapping and class_id is None:
+                    class_id = len(class_map)
+                    class_map[class_name] = class_id
+
+                if class_id is not None:
+                    xyxys.append([x0, y0, x1, y1])
+                    confidences.append(confidence / 100)
+                    class_ids.append(class_id)
+
+        if len(xyxys) == 0:
+            return Detections.empty()
+
+        return Detections(
+            xyxy=np.array(xyxys),
+            class_id=np.array(class_ids),
+            confidence=np.array(confidences),
+        )
+
     @classmethod
     def from_paddledet(cls, paddledet_result) -> Detections:
         """