Jupyter Notebook上で画像分類用のアノテーションをする

はじめに

画像のアノテーション作業、面倒くさいですよね。
特に、簡単なn値分類(n<=3)程度のアノテーションの場合、面倒なのがアノテーション自体の準備だったりすると思います。
そこで、簡単な分類タスクのアノテーション作業をJupyter Notebook上で実現することにしました。

NOTE: 本記事のコード部分は、参考記事Python: ipywidgets で Jupyter に簡単な UI を作るのコードおよびコメントを拝借し、それに修正を加えたものになります。
参考記事では、別のユースケースに適用するための説明とコードが非常に沢山かつ丁寧に含まれておりますので、興味があれば一読することをお勧めします。

コード

import random
from PIL import Image
import ipywidgets as widgets
from matplotlib import pyplot as plt

from glob import glob
from IPython.display import display


class AnnotationOnJupyter:
    def __init__(self, img_fps, labels=["犬", "猫", "鯨"]) -> None:
        """
        Args:
        - img_fps (List[str]): アノテーション対象の画像パスのリスト
        - labels (List[str]): アノテーションラベルのリスト
        """
        self.img_fps = img_fps
        self.prev_img_path = ""

        self.output = widgets.Output()
        self.ax = plt.gca()

        # コールバックを登録したボタンを用意
        btns = []
        for label in labels:
            btn = widgets.Button(description=label)
            btn.on_click(self.on_click)
            btns.append(btn)
        self.btns = btns

    def start(self):
        # boxを使うと横並びにできる。縦並びが良ければboxの行を削除して display(*btns, output) にする
        box = widgets.Box(self.btns)
        display(box, self.output)

        # セルの出力に描画されるのを抑制するために一旦アクティブな Figure を破棄する
        plt.close()

        # 最初の1回目の描画を手動でトリガーする
        blank_btn = widgets.Button(description="blank")
        blank_btn.on_click(lambda x: self.on_click(x))
        blank_btn.click()

    def on_click(self, b: widgets.Button) -> None:
        if self.prev_img_path != "":
            with open("results.txt", "a") as f:
                f.write(f"{self.prev_img_path},{b.description}\n")

        self.ax.clear()

        # NOTE: 画像パスのリストを順番にたどっていく場合はindexなどを受け取るようにすれば大丈夫
        img_path = random.choice(self.img_fps)
        self.prev_img_path = img_path
        img = Image.open(img_path)
        self.ax.imshow(img)

        #  Output に書き出す
        with self.output:
            self.output.clear_output(wait=True)
            display(self.ax.figure)


if __name__ == "__main__":
    # 画像パスのリストを適当に取得
    img_fps = glob("imgs/*.png")
    annotator = AnnotationOnJupyter(img_fps)
    annotator.start()

注意点

on_click内の処理ですが、

1. 画像Aを表示する
2. ラベルaが選択されるまで待つ
3. Aに対するラベルaを保存する

のループが1から回っているというのがわかりやすい挙動なのですが、
今回の場合は処理の最初にダミーのラベルを選択することで最初の画像を表示するところからスタートするため、選択したラベルと対応するのはひとつ前の画像パスであることに注意してください(on_clickの最初の方)。

動作確認

VSCodeのJupyter Notebookにて動作を確認しています。

さいごに

IPythonの標準ライブラリ、いつもはclear_outputくらいしか使ってなかったんですが、結構色んな事が出来そうですね。
もっといい書き方などがあればコメント等で教えていただけると嬉しいです。
読んでいただきありがとうございました。