Pillow – 画像上に物体検出の結果を描画するコード

目次

概要

Pillow を使った画像上に画像処理の物体検出の結果を矩形及びラベルで描画するためのコードについて解説します。 ディープラーニングの物体検出モデルの推論結果やOpenCV のカスケード検出器による人検出の結果の可視化に使えます。

物体検出の結果を描画する

sample1.jpg

物体検出の結果が以下の形式の dictlist で得られているとします。

  • class: ラベル
  • score: スコア
  • x1, y1: 矩形の左上の座標
  • x2, y2: 矩形の右下の座標
In [1]:
detection = [
    {
        "class": "bicycle",
        "score": 0.99,
        "x1": 103,
        "y1": 75,
        "x2": 367,
        "y2": 287,
    },
    {
        "class": "truck",
        "score": 0.92,
        "x1": 310,
        "y1": 56,
        "x2": 448,
        "y2": 110,
    },
    {
        "class": "dog",
        "score": 0.99,
        "x1": 82,
        "y1": 145,
        "x2": 204,
        "y2": 349,
    },
]

draw_boxes(img, detection, class_names) が検出結果を描画するための関数です。 img には画像を表す PIL.Image オブジェクト、detection が上記検出結果、class_names にはすべてのクラス一覧を指定します。クラス一覧は同じクラスは常に同じ色で描画するために必要となります。

In [2]:
from pathlib import Path

from matplotlib import pyplot as plt
from PIL import Image, ImageDraw, ImageFont


def draw_boxes(img, detection, class_names):
    draw = ImageDraw.Draw(img, mode="RGBA")

    # 色の一覧を作成する。
    cmap = plt.cm.get_cmap("hsv", len(class_names) + 1)

    # フォントを作成する。
    fontsize = max(15, int(0.03 * min(img.size)))
    fontname = "DejaVuSerif-Bold"
    font = ImageFont.truetype(fontname, size=fontsize)

    for bbox in detection:
        # 色を取得する。
        class_id = class_names.index(bbox["class"])
        color = cmap(class_id, bytes=True)

        # ラベル
        caption = bbox["class"]
        if "score" in bbox:
            caption += f" {bbox['score']:.0%}"  # "score" が存在する場合はパーセントで表示する。

        # 矩形を描画する。
        draw.rectangle(
            (bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]), outline=color, width=3
        )

        # ラベルを描画する。
        text_w, text_h = draw.textsize(caption, font=font)
        text_x2 = bbox["x1"] + text_w - 1
        text_y2 = bbox["y1"] + text_h - 1

        draw.rectangle((bbox["x1"], bbox["y1"], text_x2, text_y2), fill=color)
        draw.text((bbox["x1"], bbox["y1"]), caption, fill="black", font=font)
    
    return img


# 画像を読み込む。
img = Image.open("sample1.jpg")

# すべてのクラスの一覧
class_names = ["bicycle", "truck", "car", "dog"]

# 矩形を描画する。
img = draw_boxes(img, detection, class_names)

img

解説

色の一覧

N 種類の色の一覧が欲しい場合に matplotlib のカラーマップが便利です。 pyplot.cm.get_cmap(<カラーマップ名>, N) で指定したカラーマップに基づき N 種類の色が作成されます。区別がつきやすい色の一覧を作るにはカラーマップ hsv がおすすめです。hsv 色空間の場合、hue=0hue=255 は同じ色になり区別がつかなくなってしまうため、N + 1 として1個多めに色を作成しています。 カラーマップから index=i の RGB 値を取得するには、cmap(i, bytes=True) を呼び出します。

In [3]:
cmap = plt.cm.get_cmap("hsv", len(class_names) + 1)

color = cmap(1, bytes=True)  # index=1 の色を取得する。
print(color)
(133, 255, 0, 255)
クラス ID色 (RGB)
0(255, 0, 0)
1(133, 255, 0)
2(0, 255, 243)
3(109, 0, 255)
4(255, 0, 23)

テキストの描画

Pillow でテキストを描画するには、まず ImageFont.truetype(<フォント名>, size=<フォントサイズ>) でフォントオブジェクトを作成します。フォント名はその環境で利用できるものを指定します。日本語を描画したい場合は日本語フォントを指定する必要があります。 フォントサイズは画像の大きさに合わせて、調整するようにしています。テキストは text() で描画できます。

In [4]:
# フォントを作成する。
fontsize = max(15, int(0.03 * min(img.size)))
fontname = "DejaVuSerif-Bold"  # 適当なものに変える
font = ImageFont.truetype(fontname, size=fontsize)

矩形の描画

Pillow では、矩形は rectangle() で描画できます。

背景色に応じて、文字の色を変える

矩形の色によっては黒のテキストは見ずらいかもしれません。その改善策として、背景の色に応じて、黒と白を自動で切り替えるようにします。 一般に、背景の輝度値が高い場合は黒い文字が見やすく、逆に低い場合は白い文字が見やすくなります。RGB から輝度値は以下の式で計算できます。

$$ \text{brightness} = R * 0.299 + G * 0.587 + B * 0.114 $$

brightness の値が適当な閾値より大きい値の場合は明るい色ということなので、文字は黒、逆に閾値以下の場合は暗い色ということなので文字は白にします。

In [5]:
from pathlib import Path

from matplotlib import pyplot as plt
from PIL import Image, ImageDraw, ImageFont


def get_text_color(color):
    r, g, b, a = color
    brightness = r * 0.299 + g * 0.587 + b * 0.114
    return "black" if brightness > 180 else "white"


def draw_boxes(img, detection, class_names):
    draw = ImageDraw.Draw(img, mode="RGBA")

    # 色の一覧を作成する。
    cmap = plt.cm.get_cmap("hsv", len(class_names) + 1)

    # フォントを作成する。
    fontsize = max(15, int(0.03 * min(img.size)))
    fontname = "DejaVuSerif-Bold"
    font = ImageFont.truetype(fontname, size=fontsize)

    for bbox in detection:
        # 色を取得する。
        class_id = class_names.index(bbox["class"])
        color = cmap(class_id, bytes=True)

        # ラベル
        caption = bbox["class"]
        if "score" in bbox:
            caption += f" {bbox['score']:.0%}"  # "score" が存在する場合はパーセントで表示する。

        # 矩形を描画する。
        draw.rectangle(
            (bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]), outline=color, width=3
        )

        # ラベルを描画する。
        text_w, text_h = draw.textsize(caption, font=font)
        text_x2 = bbox["x1"] + text_w - 1
        text_y2 = bbox["y1"] + text_h - 1

        text_color = get_text_color(color)
        draw.rectangle((bbox["x1"], bbox["y1"], text_x2, text_y2), fill=color)
        draw.text((bbox["x1"], bbox["y1"]), caption, fill=text_color, font=font)

    return img


# 画像を読み込む。
img = Image.open("sample1.jpg")

# すべてのクラスの一覧
class_names = ["bicycle", "truck", "car", "dog"]

# 矩形を描画する。
img = draw_boxes(img, detection, class_names)

img

コメント

コメントする

目次