When working with detection models in a pipeline you often need to visualize the bounding boxes to sanity check the predictions. Bounding boxes can be defined in different ways so this is also useful to confirm the convention being used.
For example:
-
[xmin, ymin, width, height]The lower left corner of the bounding box is specified along with the height and width of the box. -
[xmin, ymin, xmax, ymax]The lower left corner and the upper right corner of the box are specified.
While visualization of bounding boxes on top of images can be done using libraries such as detectron2, in many cases it is easier to use a simple script to plot images and boxes without having to install a library or modify the input format to suit the API.
Here is a code snippet that you can copy paste and adapt to plot bounding boxes on top of images using matplotlib. If you need to add the category label annotation e.g. “chair” you can pass in an optional list of labels to the function.
from typing import List, Optional
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import patches
from pathlib import Path
def plot_bboxes(
image_file: str,
bboxes: List[List[float]],
xywh: bool = True,
labels: Optional[List[str]] = None
) -> None:
"""
Args:
image_file: str specifying the image file path
bboxes: list of bounding box annotations for all the detections
xywh: bool, if True, the bounding box annotations are specified as
[xmin, ymin, width, height]. If False the annotations are specified as
[xmin, ymin, xmax, ymax]. If you are unsure what the mode is try both
and check the saved image to see which setting gives the
correct visualization.
"""
fig = plt.figure()
# add axes to the image
ax = fig.add_axes([0, 0, 1, 1])
image_folder = Path(image_file).parent
# read and plot the image
image = plt.imread(image_file)
plt.imshow(image)
# Iterate over all the bounding boxes
for i, bbox in enumerate(bboxes):
if xywh:
xmin, ymin, w, h = bbox
else:
xmin, ymin, xmax, ymax = bbox
w = xmax - xmin
h = ymax - ymin
# add bounding boxes to the image
box = patches.Rectangle(
(xmin, ymin), w, h, edgecolor="red", facecolor="none"
)
ax.add_patch(box)
if labels is not None:
rx, ry = box.get_xy()
cx = rx + box.get_width()/2.0
cy = ry + box.get_height()/8.0
l = ax.annotate(
labels[i],
(cx, cy),
fontsize=8,
fontweight="bold",
color="white",
ha='center',
va='center'
)
l.set_bbox(
dict(facecolor='red', alpha=0.5, edgecolor='red')
)
plt.axis('off')
outfile = os.path.join(image_folder, "image_bbox.png")
fig.savefig(outfile)
print("Saved image with detections to %s" % outfile)
Example image with bounding boxes and labels
# Download to local
!wget http://farm9.staticflickr.com/8168/7101882503_99d275c5cf_z.jpg -q -O image_input.jpg
bboxes = torch.tensor(
[
[ 12.6724, 55.2408, 205.2642, 182.4455],
[321.3203, 270.7025, 395.9581, 387.7912],
[ 18.9184, 254.0521, 124.5591, 356.6135],
[257.9608, 41.3574, 390.7419, 185.9970],
[414.6888, 488.3409, 481.7668, 534.9436],
[511.3990, 232.6497, 585.9583, 363.9413],
[513.0018, 9.3525, 585.0695, 116.3015],
[ 28.3911, 392.1869, 166.8508, 549.6195],
[405.1056, 266.6370, 450.5845, 328.6919],
[235.0312, 205.1857, 414.1375, 384.2681],
[194.0130, 249.0482, 246.1905, 331.5975],
[548.2499, 400.4586, 631.2432, 494.9988],
]
dtype = torch.float32
)
labels = [
'bird',
'bird',
'bird',
'bird',
'bird',
'bird',
'bird',
'bird',
'bird',
'vase',
'bird',
'vase'
]
Plot the bounding boxes and labels onto the image:
plot_bboxes(
image_file="image_input.jpg", bboxes=bboxes,
xywh=False, labels=labels
)
 |
 |
|---|---|
Incorrect bounding boxes (with xywh=True) |
Correct bounding boxes (with xywh=False) |
You can modify the plot_bboxes function to add more annotations, for example add custom edgecolor based on the category label.