Skip to content

Commit c3cebb1

Browse files
committed
added font and score to viz
1 parent 522ad1c commit c3cebb1

File tree

3 files changed

+18
-9
lines changed

3 files changed

+18
-9
lines changed

vision_agent/fonts/__init__.py

Whitespace-only changes.

vision_agent/fonts/arial.ttf

22.2 MB
Binary file not shown.

vision_agent/image_utils.py

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Utility functions for image processing."""
22

33
import base64
4+
from importlib import resources
45
from io import BytesIO
56
from pathlib import Path
67
from typing import Dict, Tuple, Union
@@ -104,19 +105,27 @@ def overlay_bboxes(
104105

105106
color = {label: COLORS[i % len(COLORS)] for i, label in enumerate(bboxes["labels"])}
106107

107-
draw = ImageDraw.Draw(image)
108-
font = ImageFont.load_default()
109108
width, height = image.size
109+
fontsize = max(12, int(min(width, height) / 40))
110+
draw = ImageDraw.Draw(image)
111+
font = ImageFont.truetype(
112+
str(resources.files("vision_agent.fonts").joinpath("arial.ttf")), fontsize
113+
)
110114
if "bboxes" not in bboxes:
111115
return image.convert("RGB")
112116

113-
for label, box in zip(bboxes["labels"], bboxes["bboxes"]):
114-
box = [box[0] * width, box[1] * height, box[2] * width, box[3] * height]
115-
draw.rectangle(box, outline=color[label], width=3)
116-
label = f"{label}"
117-
text_box = draw.textbbox((box[0], box[1]), text=label, font=font)
118-
draw.rectangle(text_box, fill=color[label])
119-
draw.text((text_box[0], text_box[1]), label, fill="black", font=font)
117+
for label, box, scores in zip(bboxes["labels"], bboxes["bboxes"], bboxes["scores"]):
118+
box = [
119+
int(box[0] * width),
120+
int(box[1] * height),
121+
int(box[2] * width),
122+
int(box[3] * height),
123+
]
124+
draw.rectangle(box, outline=color[label], width=4)
125+
text = f"{label}: {scores:.2f}"
126+
text_box = draw.textbbox((box[0], box[1]), text=text, font=font)
127+
draw.rectangle((box[0], box[1], text_box[2], text_box[3]), fill=color[label])
128+
draw.text((box[0], box[1]), text, fill="black", font=font)
120129
return image.convert("RGB")
121130

122131

0 commit comments

Comments
 (0)