Detections loaded with sv.DetectionDataset.from_yolo annotate class IDs instead of class names with sv.LabelAnnotator #1772
Description
Search before asking
- I have searched the Supervision issues and found no similar bug report.
Bug
What's the problem?
- We usually convert the inference results into
sv.Detectionswith something like this:
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)When we use sv.LabelAnnotator with these detections, we get the following output:
sv.DetectionDataset.from_yololoads the dataset and stores the labels assv.Detections
dataset = sv.DetectionDataset.from_yolo(...)
_, _, detections = dataset[0]When we use sv.LabelAnnotator with these detections, we get the following output:
What's the expected result?
How to fix the problem?
Please see the Minimal Reproducible Example.
Environment
- Supervision 0.26.0rc3
- Ubuntu 20.04.6 LTS
- Python 3.10.15
Minimal Reproducible Example
Note: Please run the code in a notebook to use functions such as display.
Common code
import requests
from io import BytesIO
from PIL import Image
import numpy as np
import supervision as sv
from inference.models.utils import get_roboflow_model
# Create a dummy dataset
data = requests.get("https://raw.githubusercontent.com/jigsawpieces/dog-api-images/main/pitbull/dog-3981033_1280.jpg")
image = Image.open(BytesIO(data.content)).reduce(5)
label = np.random.rand(1, 5) / 10 + 0.5
label[:, 0] = 0
!mkdir -p /tmp/dummy_dataset/images
!mkdir -p /tmp/dummy_dataset/labels
image.save("/tmp/dummy_dataset/images/0.jpg")
np.savetxt("/tmp/dummy_dataset/labels/0.txt", label, fmt="%d %f %f %f %f")
with open("/tmp/dummy_dataset/dataset.yml", "w") as f:
f.write("""train: _
val: _
test: _
nc: 1
names: ["dummy"]""")Annotate detections from sv.Detections.from_inference
model = get_roboflow_model("yolov8s-640")
_, image, _ = dataset[0]
prediction = model.infer(image)[0]
detection = sv.Detections.from_inference(prediction)
annotated_image = box_annotator.annotate(image.copy(), detection)
annotated_image = label_annotator.annotate(annotated_image, detection)
display(Image.fromarray(annotated_image))
Annotate detections from sv.DetectionDataset.from_yolo
# Load as supervision dataset
dataset = sv.DetectionDataset.from_yolo("/tmp/dummy_dataset/images", "/tmp/dummy_dataset/labels", "/tmp/dummy_dataset/dataset.yml")
_, image, detection = dataset[0]
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(image.copy(), detection)
annotated_image = label_annotator.annotate(annotated_image, detection)
display(Image.fromarray(annotated_image))
How to Fix?
We need to add the data in the detections.
_, image, detection = dataset[0]
detection.data = {"class_name": np.array(['dummy'])}
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(image.copy(), detection)
annotated_image = label_annotator.annotate(annotated_image, detection)
display(Image.fromarray(annotated_image))
Additional
No response
Are you willing to submit a PR?
- Yes I'd like to help by submitting a PR!