[Python] YOLOで国会議員の居眠りを検出してみた [Google colaboratory]

はじめに

この記事は、特定の政党を支持したり批判するものではありません。

このような記事を見ました。

ベルギーの政府会議は、毎回Youtubeでライブ配信されています。ライブ配信が始まると、AIが議員とスマホを検出します。その政治家の映像はタグ付けされTwitterやInstagramのアカウントに投稿されます。

日本でも国会議員の居眠りが度々話題になるので、YOLOで検出できないか、Google Colaboratoryで実装してみました。

YOLO

一枚の画像から複数の物体を検出できる、YOLO (You Only Look Once)というアルゴリズムを利用します。検出した物体のクラスも分類可能です。YOLOは処理が高速で、動画に適用するなど、リアルタイムのアプリケーションも可能です。

YOLOの実行例

検出した物体はバウンディングボックスと呼ばれる矩形で示されています。それぞれのボックスの上の数値は、confidence score (信頼度スコア)と言い、「各バウンディングボックスの物体検出の正確さ」と「各クラスの予測確率」を表しています。

今回は比較的小さなモデルであるYOLO v5sを使います。その他にもモデルがあり、精度と計算コストなどから適宜選びます。

YOLOv5より引用

学習

衆議院インターネット審議中継からデータを集めようと思いましたが、画角の問題や、都合よく寝ている議員のデータが集められなかったので、これまでのニュース記事などから合計20枚の写真を集めました。

集めたデータをtrainとvalidationとtestに分け、trainとvalidationを使って学習します。その後、学習に用いていないtestデータで、寝ている人を検出できるか確認します。

trainとvalidation画像から、手動で寝ている人を見つけ、AIが学習する準備をします。このようなラベル付けはLabelImgを使って行いました。

デフォルトだと、save formatがPascal VOCになっていると思うので、YOLOに変えてください。txt形式で保存されます。寝てるか判断がつかない場合は、起きてると判断しました。

実装

# YOLOのclone
!git clone https://github.com/ultralytics/yolov5
!pip install -r /content/yolov5/requirements.txt

# google driveからデータセットをマウント
from google.colab import drive
drive.mount("/content/drive")
path = "/content/drive/MyDrive/PythonScripts/YOLO"

# MyDrive内のデータセットを/content/yolov5にコピー
import shutil
shutil.copytree(path+'/SleepDetection', '/content/yolov5/SleepDetection')

# 学習
# /content/yolov5内のtrain.pyの実行
# config.yamlに学習の設定を記載しておく
!python /content/yolov5/train.py --img 640 --batch 16 --epochs 200 --data /content/yolov5/SleepDetection/config.yaml --weights yolov5s.pt

# 学習の結果を可視化
%load_ext tensorboard
%tensorboard --logdir /content/yolov5/runs

今回はGoogleドライブ内にPythonScripts/YOLO/SleepDetectionというディレクトリを作成し、データセットを入れています。
SleepDetectionは次のような構成になっています。

SleepDetection
  | - config.yaml # データの情報
  | - train
  |      | - train_000.jpg # 訓練画像
  |      | - train_000.txt # ラベルデータ
  |　　　︙
  |
  | - val
  |      | - val_000.jpg # 検証画像
  |      | - val_000.txt # ラベルデータ
  |  　　︙
  |
  | - test
         | - test.jpg # テスト画像
         ︙

config.yamlの中はこのようになっています。

config.yaml

train: /content/yolov5/SleepDetection/train
val: /content/yolov5/SleepDetection/val
test: /content/yolov5/SleepDetection/test

# クラス数
nc: 1

# クラス名
names: ['sleeping']

もし多クラスの検出を行う際は、次のようにクラス名を増やしてください。

names: ['class01', 'class02', 'class03']

実際にテストデータで確認します。

# テストデータで確認
!python /content/yolov5/detect.py --source /content/yolov5/SleepDetection/test --weights /content/yolov5/runs/train/exp/weights/best.pt --conf 0.35 --name trained_exp --exist-ok --save-conf --save-txt

--save-txtを入れることで、検出したオブジェクトのラベルと座標も取り出せるようになります。検出した結果を利用して別の処理を行いたい場合はこれを入れておきましょう。

実行結果は次のディレクトリに入っています。

/content/yolov5/runs/SleepDetection/trained_exp/

--save-txtで取り出したラベルと座標、信頼度スコアは/label内に入っています。

結果

import matplotlib.pyplot as plt
import cv2

test_000 = cv2.imread('/content/yolov5/runs/detect/trained_exp/test_000.jpg',1)
test_000 = cv2.cvtColor(test_000, cv2.COLOR_BGR2RGB)
test_001 = cv2.imread('/content/yolov5/runs/detect/trained_exp/test_001.jpg',1)
test_001 = cv2.cvtColor(test_001, cv2.COLOR_BGR2RGB)

fig = plt.figure(figsize=(11,8),dpi=200)
plt.subplot(221)
plt.title('test_000')
plt.imshow(test_000)

plt.subplot(222)
plt.title('test_001')
plt.imshow(test_001)

plt.show()

出力結果

今回はたかだか20個のデータでしたが、ある程度の精度で検出できたと思います。データを増やせばもう少し精度があがりそうですね。