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

fix(pre_commit): 🎨 auto format pre-commit hooks

This commit is contained in:
pre-commit-ci[bot] 2024-07-27 22:24:39 +00:00
parent 4ce9d76b65
commit ac1fa2edca
2 changed files with 32 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
supervision/detection/core.py
View File

@ -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."

7
supervision/detection/utils.py
View File

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