画像にmatplotlibでバウンディングボックスとラベルを追加する

物体検出のテストするときによく画像の上にbounding boxとラベルを表示したくなるんだけど、 毎回同じようなコード書いててめんどくさいのでスクリプトを残しておく。

from typing import List, Tuple

import matplotlib.patches as patches
import numpy as np
from PIL import Image


def add_bboxes_to_image(ax, image: np.ndarray,
                        bboxes: List[Tuple[int, int, int, int]],
                        labels: List[str] = None,
                        label_size: int = 10,
                        line_width: int = 2,
                        border_color=(0, 1, 0, 1)) -> None:
    """
    Add bbox to ax

    :param image: dtype=np.uint8
    :param bbox: [(left, top, right, bottom)]
    :param label: List[str] or None
    :return: ax
    """
    # Display the image
    ax.imshow(image)

    if labels is None:
        labels = [None] * len(bboxes)

    for bbox, label in zip(bboxes, labels):
        # Add bounding box
        top, left, bottom, right = bbox
        rect = patches.Rectangle((left, top), right - left, bottom - top,
                                 linewidth=line_width,
                                 edgecolor=border_color,
                                 facecolor='none')
        ax.add_patch(rect)

        # label
        if label:
            bbox_props = dict(boxstyle="square,pad=0",
                              linewidth=line_width, facecolor=border_color,
                              edgecolor=border_color)
            ax.text(left, top, label,
                    ha="left", va="bottom", rotation=0,
                    size=label_size, bbox=bbox_props)
    return ax

使い方

from PIL import Image

image = np.array(Image.open('chick.jpg'))
bboxes = [(20, 130, 280, 280), (0, 0, 100, 100)]
fig, ax = plt.subplots()
add_bboxes_to_image(ax, np.uint8(image), bboxes, ['chick', '?'])

f:id:cocodrips:20200504164219p:plain

Annotations APIを使うと左上にあるラベルみたいなのを簡単に追加することができる。 matplotlib.org

今まで知らなくてboxとちょっとずらして四角とラベル追加して...とかやってた(´・ω・`) 四角以外にも矢印とかいろんなかたちでannotationを追加できる。

ひよこ画像: Serious Chick | Marji Beach | Flickr