|
1 | 1 | """Utility functions for image processing."""
|
2 | 2 |
|
3 | 3 | import base64
|
| 4 | +from importlib import resources |
4 | 5 | from io import BytesIO
|
5 | 6 | from pathlib import Path
|
6 | 7 | from typing import Dict, Tuple, Union
|
@@ -104,19 +105,27 @@ def overlay_bboxes(
|
104 | 105 |
|
105 | 106 | color = {label: COLORS[i % len(COLORS)] for i, label in enumerate(bboxes["labels"])}
|
106 | 107 |
|
107 |
| - draw = ImageDraw.Draw(image) |
108 |
| - font = ImageFont.load_default() |
109 | 108 | 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 | + ) |
110 | 114 | if "bboxes" not in bboxes:
|
111 | 115 | return image.convert("RGB")
|
112 | 116 |
|
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) |
120 | 129 | return image.convert("RGB")
|
121 | 130 |
|
122 | 131 |
|
|
0 commit comments