diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 9c7efe4e..1550f75a 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -535,17 +535,25 @@ class Detections: data=data, ) elif "segments_info" in transformers_results: - segments_info = transformers_results["segments_info"] if "segmentation" in transformers_results: - scores = np.array([segment['score'] for segment in segments_info]) - class_ids = np.array([segment['label_id'] for segment in segments_info]) - segmentation_array = transformers_results["segmentation"].cpu().detach().numpy() - masks = np.array([(segmentation_array == segment['id']).astype(bool) for segment in segments_info]) + scores = np.array([segment["score"] for segment in segments_info]) + class_ids = np.array([segment["label_id"] for segment in segments_info]) + segmentation_array = ( + transformers_results["segmentation"].cpu().detach().numpy() + ) + masks = np.array( + [ + (segmentation_array == segment["id"]).astype(bool) + for segment in segments_info + ] + ) if id2label is not None: - class_names = np.array([id2label[class_id] for class_id in class_ids]) + class_names = np.array( + [id2label[class_id] for class_id in class_ids] + ) data[CLASS_NAME_DATA_FIELD] = class_names return cls( @@ -554,15 +562,24 @@ class Detections: confidence=scores, class_id=class_ids, data=data, - ) + ) elif "png_string" in transformers_results: - class_ids = np.array([segment['category_id'] for segment in segments_info]) + class_ids = np.array( + [segment["category_id"] for segment in segments_info] + ) segmentation_array = png_to_mask(transformers_results["png_string"]) - masks = np.array([(segmentation_array == segment['id']).astype(bool) for segment in segments_info]) + masks = np.array( + [ + (segmentation_array == segment["id"]).astype(bool) + for segment in segments_info + ] + ) if id2label is not None: - class_names = np.array([id2label[class_id] for class_id in class_ids]) + class_names = np.array( + [id2label[class_id] for class_id in class_ids] + ) data[CLASS_NAME_DATA_FIELD] = class_names return cls( @@ -570,7 +587,7 @@ class Detections: mask=masks, class_id=class_ids, data=data, - ) + ) else: raise NotImplementedError( "Only object detection and semantic segmentation results are supported." diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 0884d0d5..00b3693a 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1,10 +1,10 @@ +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 -import io from PIL import Image from supervision.config import CLASS_NAME_DATA_FIELD @@ -1003,6 +1003,7 @@ 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. @@ -1015,6 +1016,6 @@ def png_to_mask(png_string): of the image. """ image = Image.open(io.BytesIO(png_string)) - mask = np.array(image, dtype= np.uint8) + mask = np.array(image, dtype=np.uint8) - return mask[:,:,0] + return mask[:, :, 0]