YOLOv5をGoogleColaboratoryで動かして静止画と動画の物体検出をしてみた

背景

会社の友人向けにハンズオンでのAI体験会を開催することになりました。画像データを扱うことが視覚的に面白いと思ったのと、環境構築で苦労したくなかったので、YOLOv５をGoogleColaboratoryで動かして、静止画と動画の物体検出にトライしてみました。

参考サイト

以下のサイトを参考にさせていただきました。
https://laboratory.kazuuu.net/i-used-yolo-v5-with-google-colaboratory/
https://qiita.com/shoku-pan/items/31bf3c975b73db153121

解析ファイルの準備

まず、画像と動画を準備して、それぞれtest.jpg、test.mp4として保存します。ファイル形式はソースコードを修正すれば他の形式でも問題ありませんが、今回はこのファイル形式で進めます。
なお、今回は下記手法・サイトで見つけた著作権フリーの画像で実行しています。

画像：https://ferret-plus.com/337
動画：https://pixabay.com/ja/videos/

次に、GoogleDriveにフォルダ構成を作成します。GoogleColaboratoryからGoogleDriveにアクセスして、Drive上に保存されている画像および動画の物体検出を行います。その準備として、下図のようにMyDrive直下にYOLOフォルダを作成し、そこに上記test.jpgとtest.mp4を格納します。

GoogleColaboratoryの準備

次に、GoogleCoraboratoryを使えるようにします。先ほど作成したYOLOフォルダ内の何もないところで右クリックをして、「その他」→「Google Colaboratory」を選択すると、Jupyter Notebook形式のGoogleColaboratoryが開きます。簡単ですね！

GoogleColaboratoryへの実装

YOLOv5の準備

空のセルに、下記の区切りでコードを入力し、一つずつ実行していきます（セル内でShift+Enterで実行できますね）

#GoogleDriveへのアクセスをできるようにする
from google.colab import drive
drive.mount('/content/drive')

上記を実行すると、GoogleDriveへの認証を求められるので、許可してください

#YOLO5のインストール
!git clone https://github.com/ultralytics/yolov5
!pip install -r yolov5/requirements.txt

これでYOLOv5がGoogleColaboratory上で使えるようになります。

#フォルダを移動
!cd /content/yolov5

以上で環境構築終了です。簡単ですねー！

画像の解析

ここまで来れば、あとはファイルを指定してYOLOを実行するだけです。

#画像解析
!python /content/yolov5/detect.py --source /content/drive/MyDrive/YOLO/test.jpg --weights yolov5x.pt --conf 0.5 --exist-ok --line-thickness 3

参考までに、今回使用している設定値の説明は以下の通りです。
--source：YOLOにかけるファイルを指定。
--weights：YOLOv5の中のどのモデルを使用するかを指定。今回は最も精度の高いyolov5x.ptを使用しています。そのほかにも、以下のようなモデルがあるので、精度と速度のバランスを見て選択できます。（参考：https://github.com/ultralytics/yolov5）

--conf：物体を検出する閾値の設定。
--exist-ok：既存ファイルへの上書き保存許可。
--line-thickness：解析画像の識別結果の文字の大きさと枠の太さ。解析結果を見ながら設定するのがよいです。

#解析結果をGoogleColaboratory上に表示する
import cv2
im =cv2.imread('/content/drive/MyDrive/YOLO/test.jpg')#元画像
im2 =cv2.imread('/content/yolov5/runs/detect/exp/test.jpg')#解析画像

# 元画像の表示
from matplotlib import pyplot as plt
plt.figure(figsize=(12,8))
plt.axis('off')
im_rab = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
plt.imshow(im_rab)

# 解析画像の表示
from matplotlib import pyplot as plt
plt.figure(figsize=(12,8))
plt.axis('off')
im2_rab = cv2.cvtColor(im2, cv2.COLOR_BGR2RGB)
plt.imshow(im2_rab)

上記を実行すると、以下のように元画像と解析画像が表示されます。

＜元画像＞

＜解析画像＞

YOLOが実行されて、人が検出できていることがわかります！右手前の人は、リュックのベルトだけしかみえていないのに、backpackと認識されていますね。YOLO恐るべし。。。

YOLOの解析結果は、GoogleColaboratoryにインストールしたyolov5フォルダに格納されるようです。GoogleColabがリセットされる（≒閉じる）とそのフォルダは消えてしまうので、解析結果を最初にGoogleDriveに作成したYOLOフォルダにtest_yolo.jpgとして移動・格納します。

import os
import shutil

# 解析結果のファイルパス
path1 = '/content/yolov5/runs/detect/exp/test.jpg'
# 解析結果をコピーした後のファイルパス
path2 = '/content/yolov5/runs/detect/exp/test_yolo.jpg'
# 格納先のファイルパス
path3 = '/content/drive/MyDrive/YOLO/test_yolo.jpg'

# ファイルの存在確認
if os.path.exists(path3) == True: #path3のファイルがあるかどうかを確認する
  print("すでに解析ファイル(test_yolo.jpg)がGoogleDriveに保存されています")

elif os.path.exists(path1) == True: #path1のファイルがあるかどうかを確認する
  # ファイル名のコピー
  shutil.copy(path1, path2) #解析ファイルの名前を変えてコピー
  shutil.move('/content/yolov5/runs/detect/exp/test_yolo.jpg', '/content/drive/MyDrive/YOLO') #GoogleDriveへ移動
  print("ファイル名の変更と　GoogleDriveへの移動が終わりました（ファイル名：test_yolo.jpg）")

else:
  print("YOLO５の解析が未実行です。")

上記を実行して
ファイル名の変更と　GoogleDriveへの移動が終わりました（ファイル名：test_yolo.jpg）
と表示されれば、GoogleDriveのYOLOフォルダにtest_yolo.jpgが保存されているはずです。（少し遅れて保存されます）

動画の物体検出

次に、動画の物体検出を行います。動画の解析も画像の時とほぼ同じです。動画は時間がかかるので、ここでは高速なyolov5n.ptを使用しています。

#動画解析
!python /content/yolov5/detect.py --source /content/drive/MyDrive/YOLO/test.mp4 --weights yolov5n.pt --conf 0.3 --exist-ok

先ほどと同じように、GoogleDriveのYOLOフォルダにファイルを移動・格納します。

import os
import shutil

# 解析結果のファイルパス
path1 = '/content/yolov5/runs/detect/exp/test.mp4'
# 解析結果をコピーした後のファイルパス
path2 = '/content/yolov5/runs/detect/exp/test_yolo.mp4'
# 格納先のファイルパス
path3 = '/content/drive/MyDrive/YOLO/test_yolo.mp4'

# ファイルの存在確認
if os.path.exists(path3) == True: #path3のファイルがあるかどうかを確認する
  print("すでに解析ファイル(test_yolo.mp4)がGoogleDriveに保存されています")

elif os.path.exists(path1) == True: #path1のファイルがあるかどうかを確認する
  # ファイル名のコピー
  shutil.copy(path1, path2) #解析ファイルの名前を変えてコピー
  shutil.move('/content/yolov5/runs/detect/exp/test_yolo.mp4', '/content/drive/MyDrive/YOLO') #GoogleDriveへ移動
  print("ファイル名の変更と　GoogleDriveへの移動が終わりました（ファイル名：test_yolo.mp4）")

else:
  print("YOLO５の解析が未実行です。")

YOLOフォルダに格納したtest_yolo.mp4を開いてみると、、、

動画でもしっかり物体検出ができています！

まとめ

今回はGoogleColaboratory上でYOLOを実装してみました。環境構築も簡単で、他の人と一緒に触ったり、簡単な物体検出をする分にはとてもいい手段と感じました。
一方で、外付けwebカメラなどをつないでリアルタイム検出することができなさそうなので、その場合はPCなどにYOLOv5をインストールして使う必要がありそうです。こちらも余裕があればやってみたいと思います。