物体検出のテストするときによく画像の上に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', '?'])
Annotations APIを使うと左上にあるラベルみたいなのを簡単に追加することができる。 matplotlib.org
今まで知らなくてboxとちょっとずらして四角とラベル追加して...とかやってた(´・ω・｀) 四角以外にも矢印とかいろんなかたちでannotationを追加できる。
