jupyter lab上で画像のラベリングをする

やりたいこと

AIに使いたい画像のラベリングをしたいが、GUIを作るのが面倒。。。と思い、
JupyterLab上でできないかーと考え、色々調べてたらできたのでそのときのメモ。
matplotで表示した画像を更新しながらinputでラベリングを入力するようなイメージ。

実行環境

• Python 3.8.5
• jupyter lab 2.2.2
• mac(windows10でも実行済み)

コード

sample_imgフォルダにある画像をラベリングする。

準備

import cv2
import matplotlib.pyplot as plt
import os
import glob
import IPython
# 画像のpathを取得
img_list = glob.glob(os.path.join(r"sample_img","*.jpg"))
img_list
>>> ['sample_img/img3.jpg', 'sample_img/img2.jpg', 'sample_img/img1.jpg']

画像を表示するコード

name_list = []
for img_path in img_list:
    # 画像の読み込み
    img = cv2.imread(img_path)
    # 画像を表示するフレーム
    fig = plt.figure(figsize=(5,5))
    ax = fig.add_subplot(1,1,1)
    ax.imshow(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
    plt.title(img_path)
    plt.pause(.01)
    
    # ラベル付けするためのテキストボックス
    comment = input()
   
    if comment == "break":
        break
    else:
        # 表示している画像をクリア
        IPython.display.clear_output()
        # Inputに入力した文字列をリストに追加
        name_list.append(comment)

動作画面はこんな感じ。
pass

確認

# name_listの確認
name_list
>>> ['dog', 'cat', 'cat']

最後に

Input部分をラジオボタンみたいにできたらいいな〜.