xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cyclic dependency fix + code refactor

This commit is contained in:
Shadab 2024-08-03 00:36:36 +01:00
parent a5ae72d7ae
commit da261ff8cf
2 changed files with 127 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
supervision/detection/core.py
View File

@ -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:

129
supervision/detection/tools/transformers.py
View File

@ -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,
)