diff --git a/docs/how_to/detect_and_annotate.md b/docs/how_to/detect_and_annotate.md index 7bc2f970..cdd13030 100644 --- a/docs/how_to/detect_and_annotate.md +++ b/docs/how_to/detect_and_annotate.md @@ -102,7 +102,7 @@ Now that we have predictions from a model, we can load them into Supervision. We can do so using the [`sv.Detections.from_transformers`](detection/core/#supervision.detection.core.Detections.from_transformers) method, which accepts model results from both detection and segmentation models. - ```{ .py hl_lines="2 19" } + ```{ .py hl_lines="2 19-21" } import torch import supervision as sv from PIL import Image @@ -121,7 +121,9 @@ Now that we have predictions from a model, we can load them into Supervision. target_size = torch.tensor([[height, width]]) results = processor.post_process_object_detection( outputs=outputs, target_sizes=target_size)[0] - detections = sv.Detections.from_transformers(results) + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label) ``` You can load predictions from other computer vision frameworks and libraries using: @@ -180,7 +182,7 @@ Finally, we can annotate the image with the predictions. Since we are working wi === "Transformers" - ```{ .py hl_lines="21-27" } + ```{ .py hl_lines="23-30" } import torch import supervision as sv from PIL import Image @@ -199,7 +201,9 @@ Finally, we can annotate the image with the predictions. Since we are working wi target_size = torch.tensor([[height, width]]) results = processor.post_process_object_detection( outputs=outputs, target_sizes=target_size)[0] - detections = sv.Detections.from_transformers(results) + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label) bounding_box_annotator = sv.BoundingBoxAnnotator() label_annotator = sv.LabelAnnotator() @@ -274,7 +278,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me === "Transformers" - ```{ .py hl_lines="24-28 33" } + ```{ .py hl_lines="26-30 35" } import torch import supervision as sv from PIL import Image @@ -293,15 +297,17 @@ override this behavior by passing a list of custom `labels` to the `annotate` me target_size = torch.tensor([[height, width]]) results = processor.post_process_object_detection( outputs=outputs, target_sizes=target_size)[0] - detections = sv.Detections.from_transformers(results) + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label) bounding_box_annotator = sv.BoundingBoxAnnotator() label_annotator = sv.LabelAnnotator() labels = [ - f"{model.config.id2label[class_id]} {confidence:.2f}" - for class_id, confidence - in zip(detections.class_id, detections.confidence) + f"{class_name} {confidence:.2f}" + for class_name, confidence + in zip(detections['class_name'], detections.confidence) ] annotated_image = bounding_box_annotator.annotate( @@ -383,15 +389,17 @@ that will allow you to draw masks instead of boxes. target_size = torch.tensor([[height, width]]) results = processor.post_process_segmentation( outputs=outputs, target_sizes=target_size)[0] - detections = sv.Detections.from_transformers(results) + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label) mask_annotator = sv.MaskAnnotator() label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER) labels = [ - f"{model.config.id2label[class_id]} {confidence:.2f}" - for class_id, confidence - in zip(detections.class_id, detections.confidence) + f"{class_name} {confidence:.2f}" + for class_name, confidence + in zip(detections['class_name'], detections.confidence) ] annotated_image = mask_annotator.annotate( diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 554495ab..016426b4 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -189,7 +189,7 @@ class Detections: Args: ultralytics_results (ultralytics.yolo.engine.results.Results): - The output Results instance from YOLOv8 + The output Results instance from Ultralytics Returns: Detections: A new Detections object. @@ -394,12 +394,19 @@ class Detections: ) -> Detections: """ Creates a Detections instance from object detection or segmentation - [transformer](https://github.com/huggingface/transformers) inference result. + [Transformer](https://github.com/huggingface/transformers) inference result. + + !!! note + + Class names can be accessed using the key `class_name` in the returned + object's data attribute. Args: - transformers_results (dict): A dictionary containing the scores, labels, and boxes for an image - as predicted by the Transformers model. - id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to class names. If None, class names will not be included. + transformers_results (dict): The output of Transformers model inference. A + dictionary containing the `scores`, `labels`, `boxes` and `masks` keys. + id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to + class names. If provided, the resulting Detections object will contain + `class_name` data field with the class names. Returns: Detections: A new Detections object. @@ -424,7 +431,11 @@ class Detections: target_size = torch.tensor([[height, width]]) results = processor.post_process_object_detection( outputs=outputs, target_sizes=target_size)[0] - detections = sv.Detections.from_transformers(results) + + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label + ) ``` """ # noqa: E501 // docs @@ -507,7 +518,7 @@ class Detections: !!! note - Class names can be accessed using the key 'class_name' in the returned + Class names can be accessed using the key `class_name` in the returned object's data attribute. Args: