Spaces:
Running
Running
File size: 1,951 Bytes
111afa2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import supervision as sv
from smolagents import AgentImage, Tool
class LabelAnnotatorTool(Tool):
name = "label_annotator"
description = """
Given an image and a list of detections, draw labels on the image.
The image is a PIL image.
The detections are an object of type supervision.Detections. You can use the task_inference_output_converter tool to obtain the proper format for the detections.
The output is the image with the labels drawn on it.
"""
inputs = {
"image": {
"type": "image",
"description": "The image to annotate the labels on.",
},
"detections": {
"type": "object",
"description": """
The detections to annotate on the image.
The detections are an object of type supervision.Detections.
You can use the task_inference_output_converter tool to obtain the proper format for the detections.
""",
},
"text_position": {
"type": "string",
"description": """
The position of the label relative to the bounding box.
The supported positions are:
- CENTER
- CENTER_LEFT
- CENTER_RIGHT
- TOP_CENTER
- TOP_LEFT
- TOP_RIGHT
- BOTTOM_LEFT
- BOTTOM_CENTER
- BOTTOM_RIGHT
- CENTER_OF_MASS
""",
},
}
output_type = "image"
def __init__(self):
super().__init__()
def forward(
self,
image: AgentImage,
detections: sv.Detections,
text_position: str,
):
label_annotator = sv.LabelAnnotator(text_position=sv.Position(text_position))
annotated_image = label_annotator.annotate(
scene=image, detections=detections, labels=detections.metadata["labels"]
)
return annotated_image
|