Detections loaded with sv.DetectionDataset.from_yolo annotate class IDs instead of class names with sv.LabelAnnotator

[patel-zeel]

Search before asking

• I have searched the Supervision issues and found no similar bug report.

Bug

What's the problem?

1. We usually convert the inference results into sv.Detections with 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:

2. sv.DetectionDataset.from_yolo loads the dataset and stores the labels as sv.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!

[SkalskiP]

Hi @patel-zeel 👋🏻 Another great find! This is indeed a bug. It looks like when loading labels using sv.DetectionDataset.from_yolo, we are not saving the class_names key in sv.Detections.data.

This is how sv.LabelAnnotator works. First, it checks if sv.Detections has the key class_names. If so, we print the label found in it; if not, we use sv.Detections.class_id as a fallback.

[patel-zeel]

Thank you for the confirmation, @SkalskiP! I checked that sv.DetectionDataset.from_coco and sv.DetectionDataset.from_pascal_voc also have the same issue. In this case, would the following location be optimal to fix this bug?

supervision/supervision/dataset/core.py

Lines 84 to 87 in 1b5c592

	self.annotations = annotations
	# Eliminate duplicates while preserving order
	self.image_paths = list(dict.fromkeys(images))

Proposed changes:

self.annotations = annotations
np_classes = np.array(self.classes)
for annotation in self.annotations:
    annotation.data["class_name"] = np_classes[annotation.class_id]

# Eliminate duplicates while preserving order
self.image_paths = list(dict.fromkeys(images))