From da261ff8cfd0b93f31f52f765b970ec787a1a76c Mon Sep 17 00:00:00 2001 From: Shadab Date: Sat, 3 Aug 2024 00:36:36 +0100 Subject: [PATCH] cyclic dependency fix + code refactor --- supervision/detection/core.py | 39 ++++-- supervision/detection/tools/transformers.py | 129 +++++++++++++++----- 2 files changed, 127 insertions(+), 41 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index e2843994..3ea3b146 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -442,7 +442,8 @@ class Detections: cls, transformers_results: dict, id2label: Optional[Dict[int, str]] = None ) -> Detections: """ - Creates a Detections instance from object detection or segmentation + Creates a Detections instance from object detection or panoptic, semantic + and instance segmentation [Transformer](https://github.com/huggingface/transformers) inference result. Args: @@ -490,17 +491,37 @@ class Detections: Class names values can be accessed using `detections["class_name"]`. """ # noqa: E501 // docs - if "boxes" in transformers_results: - return process_detection_result(transformers_results, id2label) + 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['mask'], + class_id=processed_result['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 + )) + if "masks" in transformers_results or "png_string" in transformers_results: - return process_transformers_v4_segmentation_result( + return convert_to_detections(process_transformers_v4_segmentation_result( transformers_results, id2label - ) - else: - return process_transformers_v5_segmentation_result( - transformers_results, id2label - ) + )) + + if "boxes" in transformers_results: + return convert_to_detections(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 4224556d..cc44e4ef 100644 --- a/supervision/detection/tools/transformers.py +++ b/supervision/detection/tools/transformers.py @@ -3,12 +3,20 @@ from typing import Dict, Optional import numpy as np from supervision.config import CLASS_NAME_DATA_FIELD -from supervision.detection.core import Detections from supervision.detection.utils import mask_to_xyxy, png_to_mask def get_data(class_ids: np.ndarray, id2label: Optional[Dict[int, str]]) -> dict: - """Helper function to create data dictionary with class names if available.""" + """ + Helper function to create 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. + + Returns: + dict: Dictionary containing class names if id2label is provided. + """ data = {} if id2label is not None: class_names = np.array([id2label[class_id] for class_id in class_ids]) @@ -18,8 +26,19 @@ def get_data(class_ids: np.ndarray, id2label: Optional[Dict[int, str]]) -> dict: def process_tensor_result( segmentation_array: np.ndarray, id2label: Optional[Dict[int, str]] -) -> Detections: - """Process segmentation array result for segmentation.""" +) -> dict: + """ + Helper function to process result of transformers function + post_process_panoptic_segmentation. + + Args: + segmentation_array (np.ndarray): Segmentation array. + id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. + + Returns: + dict: Processed segmentation result including bounding boxes, masks, + class IDs, and data. + """ class_ids = np.unique(segmentation_array) masks = np.stack( [(segmentation_array == class_id).astype(bool) for class_id in class_ids], @@ -27,19 +46,31 @@ def process_tensor_result( ) data = get_data(class_ids, id2label) - return Detections( + return dict( xyxy=mask_to_xyxy(masks), mask=masks, class_id=class_ids, data=data ) def process_detection_result( detection_result: dict, id2label: Optional[Dict[int, str]] -) -> Detections: - """Process detection results containing boxes and labels.""" +) -> dict: + """ + Helper function to process result of transformers functions + post_process_object_detection and post_proces. + + Args: + detection_result (dict): Dictionary containing detection results with keys + 'boxes', 'labels', and 'scores'. + id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. + + Returns: + dict: Processed detection result including bounding boxes, confidence scores, + class IDs, and data. + """ class_ids = detection_result["labels"].cpu().detach().numpy().astype(int) data = get_data(class_ids, id2label) - return Detections( + return dict( xyxy=detection_result["boxes"].cpu().detach().numpy(), confidence=detection_result["scores"].cpu().detach().numpy(), class_id=class_ids, @@ -49,27 +80,33 @@ def process_detection_result( def process_transformers_v4_segmentation_result( segmentation_result: dict, id2label: Optional[Dict[int, str]] -) -> Detections: +) -> dict: """ - Process Transformers v4 segmentation results. + Helper function to process result of transformers functions + post_process_panoptic, post_process_segmentation and post_process_instance. Args: - segmentation_result (dict): Dictionary containing segmentation results with keys 'masks', 'labels', and 'scores'. + segmentation_result (dict): Dictionary containing segmentation results with keys + 'masks', 'labels', and 'scores'. id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. Returns: - Detections: A Detections object created from the segmentation results. + dict: Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. """ - if "png_string" in segmentation_result: return process_png_segmentation_result(segmentation_result, id2label) else: + boxes = None + if "boxes" in segmentation_result: + boxes = segmentation_result["boxes"].cpu().detach().numpy() + masks = segmentation_result["masks"].cpu().detach().numpy().astype(bool) class_ids = segmentation_result["labels"].cpu().detach().numpy().astype(int) - return Detections( - xyxy=mask_to_xyxy(masks), - mask=masks, + return dict( + xyxy=boxes if boxes is not None else mask_to_xyxy(masks), + 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), @@ -78,29 +115,44 @@ def process_transformers_v4_segmentation_result( def process_transformers_v5_segmentation_result( segmentation_result: dict, id2label: Optional[Dict[int, str]] -) -> Detections: +) -> dict: """ - Process Transformers v5 segmentation results. + Helper function to process result of transformers functions + post_process_semantic_segmentation, post_process_instance_segmentation and + post_process_panoptic_segmentation. Args: - segmentation_result (Union[dict, np.ndarray]): Either a dictionary containing segmentation results or an ndarray representing a segmentation map. + segmentation_result (Union[dict, np.ndarray]): Either a dictionary containing + segmentation results or an ndarray representing a segmentation map. id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. Returns: - Detections: A Detections object created from the segmentation results. + dict: Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. """ - if segmentation_result.__class__.__name__ == "Tensor": segmentation_array = segmentation_result.cpu().detach().numpy() return process_tensor_result(segmentation_array, id2label) - return process_png_segmentation_result(segmentation_result, id2label) + return process_segmentation_result(segmentation_result, id2label) def process_segmentation_result( segmentation_result: dict, id2label: Optional[Dict[int, str]] -) -> Detections: - """Process segmentation results with masks and scores.""" +) -> dict: + """ + Helper function to process result of transformers functions + post_process_semantic_segmentation and post_process_instance_segmentation. + + Args: + segmentation_result (dict): Dictionary containing segmentation results with keys + 'segments_info' and 'segmentation'. + id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. + + Returns: + dict: Processed segmentation result including bounding boxes, masks, confidence + scores, class IDs, and data. + """ segments_info = segmentation_result["segments_info"] scores = np.array([segment["score"] for segment in segments_info]) class_ids = np.array([segment["label_id"] for segment in segments_info]) @@ -113,7 +165,7 @@ def process_segmentation_result( ) data = get_data(class_ids, id2label) - return Detections( + return dict( xyxy=mask_to_xyxy(masks), mask=masks, confidence=scores, @@ -123,12 +175,22 @@ def process_segmentation_result( def process_png_segmentation_result( - png_result: dict, id2label: Optional[Dict[int, str]] -) -> Detections: - """Process segmentation results from a PNG string.""" - segments_info = png_result["segments_info"] + segmentation_result: dict, id2label: Optional[Dict[int, str]] +) -> dict: + """ + Helper function to process result of transformers function post_process_panoptic. + + Args: + segmentation_result (dict): Dictionary containing PNG string and segment information. + id2label (Optional[Dict[int, str]]): Dictionary mapping class IDs to class names. + + Returns: + dict: Processed segmentation result including bounding boxes, masks, + class IDs, and data. + """ + segments_info = segmentation_result["segments_info"] class_ids = np.array([segment["category_id"] for segment in segments_info]) - segmentation_array = png_to_mask(png_result["png_string"]) + segmentation_array = png_to_mask(segmentation_result["png_string"]) masks = np.array( [ (segmentation_array == segment["id"]).astype(bool) @@ -137,6 +199,9 @@ def process_png_segmentation_result( ) data = get_data(class_ids, id2label) - return Detections( - xyxy=mask_to_xyxy(masks), mask=masks, class_id=class_ids, data=data + return dict( + xyxy=mask_to_xyxy(masks), + mask=masks, + class_id=class_ids, + data=data, )