diff --git a/supervision/detection/core.py b/supervision/detection/core.py index fc0317fd..0ae3fa65 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -491,46 +491,19 @@ class Detections: Class names values can be accessed using `detections["class_name"]`. """ # noqa: E501 // docs - def convert_to_detections(processed_result: dict) -> Detections: - """ - Convert a processed result dictionary into a Detections object. - - Args: - processed_result (dict): The processed result dictionary containing the keys - 'xyxy', 'mask', 'class_id', 'data', and 'confidence'. - - Returns: - Detections: A Detections object created from the processed result. - """ - return Detections( - xyxy=processed_result["xyxy"], - mask=processed_result.get("mask"), - class_id=processed_result.get("class_id"), - data=processed_result["data"], - confidence=processed_result.get("confidence"), - ) - if ( transformers_results.__class__.__name__ == "Tensor" or "segmentation" in transformers_results ): - return convert_to_detections( - process_transformers_v5_segmentation_result( - transformers_results, id2label - ) - ) + return cls(**process_transformers_v5_segmentation_result( + transformers_results, id2label)) if "masks" in transformers_results or "png_string" in transformers_results: - return convert_to_detections( - process_transformers_v4_segmentation_result( - transformers_results, id2label - ) - ) + return cls(**process_transformers_v4_segmentation_result( + transformers_results, id2label)) if "boxes" in transformers_results: - return convert_to_detections( - process_detection_result(transformers_results, id2label) - ) + return cls(**process_detection_result(transformers_results, id2label)) @classmethod def from_detectron2(cls, detectron2_results) -> Detections: diff --git a/supervision/detection/tools/transformers.py b/supervision/detection/tools/transformers.py index 3c380167..0833e092 100644 --- a/supervision/detection/tools/transformers.py +++ b/supervision/detection/tools/transformers.py @@ -1,26 +1,38 @@ -from typing import Dict, Optional +import io +from typing import Dict, Optional, Any import numpy as np +from PIL import Image from supervision.config import CLASS_NAME_DATA_FIELD -from supervision.detection.utils import mask_to_xyxy, png_to_mask +from supervision.detection.utils import mask_to_xyxy -def get_data(class_ids: np.ndarray, id2label: Optional[Dict[int, str]]) -> dict: +def append_class_names_to_data( + class_ids: np.ndarray, + id2label: Optional[Dict[int, str]], + data: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: """ - Helper function to create data dictionary with class names if available. + Helper function to create or append to a data dictionary with class names if + available. Args: class_ids (np.ndarray): Array of class IDs. - id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. + id2label (Optional[Dict[int, str]]): A map from index to label. Typically part + of `transformers` model configuration. + data (Optional[Dict[str, Any]]): An existing data dictionary to append to. Returns: - dict: Dictionary containing class names if id2label is provided. + Dict[str, Any]: Dictionary containing class names if id2label is provided. """ - data = {} + if data is None: + data = {} + if id2label is not None: class_names = np.array([id2label[class_id] for class_id in class_ids]) data[CLASS_NAME_DATA_FIELD] = class_names + return data @@ -44,7 +56,7 @@ def process_tensor_result( [(segmentation_array == class_id).astype(bool) for class_id in class_ids], axis=0, ) - data = get_data(class_ids, id2label) + data = append_class_names_to_data(class_ids, id2label, {}) return dict(xyxy=mask_to_xyxy(masks), mask=masks, class_id=class_ids, data=data) @@ -66,7 +78,7 @@ def process_detection_result( class IDs, and data. """ class_ids = detection_result["labels"].cpu().detach().numpy().astype(int) - data = get_data(class_ids, id2label) + data = append_class_names_to_data(class_ids, id2label, {}) return dict( xyxy=detection_result["boxes"].cpu().detach().numpy(), @@ -107,7 +119,7 @@ def process_transformers_v4_segmentation_result( mask=np.squeeze(masks, axis=1) if boxes is not None else masks, confidence=segmentation_result["scores"].cpu().detach().numpy(), class_id=class_ids, - data=get_data(class_ids, id2label), + data=append_class_names_to_data(class_ids, id2label ,{}), ) @@ -161,7 +173,7 @@ def process_segmentation_result( for segment in segments_info ] ) - data = get_data(class_ids, id2label) + data = append_class_names_to_data(class_ids, id2label, {}) return dict( xyxy=mask_to_xyxy(masks), @@ -188,14 +200,14 @@ def process_png_segmentation_result( """ segments_info = segmentation_result["segments_info"] class_ids = np.array([segment["category_id"] for segment in segments_info]) - segmentation_array = png_to_mask(segmentation_result["png_string"]) + label_mask = png_string_to_label_mask(segmentation_result["png_string"]) masks = np.array( [ - (segmentation_array == segment["id"]).astype(bool) + (label_mask == segment["id"]).astype(bool) for segment in segments_info ] ) - data = get_data(class_ids, id2label) + data = append_class_names_to_data(class_ids, id2label, {}) return dict( xyxy=mask_to_xyxy(masks), @@ -203,3 +215,20 @@ def process_png_segmentation_result( class_id=class_ids, data=data, ) + + +def png_string_to_label_mask(png_string: bytes) -> np.ndarray: + """ + Convert a PNG byte string to a label mask array. + + Args: + png_string (bytes): A byte string representing the PNG image. + + Returns: + np.ndarray: A label mask array with shape (H, W), where H and W + are the height and width of the image. Each unique value in the array + represents a different object or category. + """ + image = Image.open(io.BytesIO(png_string)) + mask = np.array(image, dtype=np.uint8) + return mask[:, :, 0] diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 9d33f76c..a02a6a1f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,11 +1,9 @@ -import io from itertools import chain from typing import Dict, List, Optional, Tuple, Union import cv2 import numpy as np import numpy.typing as npt -from PIL import Image from supervision.config import CLASS_NAME_DATA_FIELD from supervision.geometry.core import Vector @@ -1002,18 +1000,3 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray: ) vector_start = np.array([vector.start.x, vector.start.y]) return np.cross(vector_at_zero, anchors - vector_start) - - -def png_to_mask(png_string): - """ - Convert a PNG byte string to a binary mask array. - Args: - png_string (bytes): A byte string representing the PNG image. - - Returns: - np.ndarray: A binary mask array with shape (H, W), where H and W - are the height and width of the image. - """ - image = Image.open(io.BytesIO(png_string)) - mask = np.array(image, dtype=np.uint8) - return mask[:, :, 0]