cedro-blog

PyTorch 学習済みモデルでサクッと物体検出をしてみる

今回は、SSDの学習済みモデルを使って物体検出をやってみます。

こんにちは cedro です。

以前、YOLOv3 で物体検出をやってみましたが、PyTorchでももちろんできます。

PyTorchでは、YOLOv3と同様に、バウンディングボックスの検出とクラス分類を平行して行うことで、高速な物体検出を実現したSSD(Single Shot multibox Detection)というモデルが使えます。

ということで、今回は、SSDの学習済みモデルを使って物体検出をやってみます。

 

SSDとは

YOLOのネットワーク

YOLOは、画像を7×7のグリッドに分割し、グリッド毎に物体検出とクラス分類を平行して行い、この結果を総合して適切なバウンディングボックスを求め、そのクラス分類結果を出力します。

 

   

(引用:SSD: Single Shot MultiBox Detector:https://arxiv.org/pdf/1512.02325.pdf)

一方、SSDは、8732個のバウンディングボックスをあらかじめ設定し、バウンディングボックス毎に物体検出とクラス分類を平行して行い、最後に適切なバウンディングボックスを選択しそのクラス分類結果を出力します。

 

SSD300の学習済みモデルが物体検出できるクラス

今回、使うSSD300の学習済みモデルが物体検出できるクラスは、この20種類です。

 

コードを書きます

今回は、学習済みモデルを使いますので、実行環境はノートパソコンで十分です。もし、PyTorch (Ver1.0)がまだインストール済みでなければ、PyTorchのホームページの説明に従ってインストールして下さい。

まず、Github からサンプルコードをクローンあるいはダウンロードします。今回使用するのは、Chapter7です。

 

Chapter7の中身

Chapter7の中身です。ここに、demo2 フォルダーを追加します。今回作成するコードは、全てこのdemo2フォルダーの中に書きます。

また、SSDの学習済みモデル( ssd300_mAP_77.43_v2.pth )をダウンロードし、weights フォルダーに格納しておきます。

  

  

SSDの学習済みモデルを使って、物体検出するコードです。学習済みモデルを読み込んだら、image_dir に格納されている画像(png)を1つづつ取り出し、画像を300×300にリサイズして物体検出し、バウンディングボックスとクラス名を書き込み、detect_img に画像(png)を保存します。

学習済みモデルの実行なので、あえてGPUは使わない仕様にしています。

 

 

先程のコードを ssd_model.py で保存します。そして、image_dir フォルダー(物体検知をしたい画像を入れる)とdetect_img フォルダー(物体検知の結果画像が出力される)を作成します。

 

  

cat_girl.png

この画像 ( cat_girl.png ) をimage_dir フォルダーに格納します。

 

 

0001.png

そして、ssd_model.py を実行すると、物体検出の結果画像 (0001.png) が、detect_img フォルダーに格納されます。

 

動画を物体検出してみる

物体検出と言えば、リアルタイムに物体検出している動画(バウンディングボックスとクラス名が連続的に表示されている)を見ることがあると思いますが、今回はあれを疑似的に再現したいと思います。

  

   

動画から静止画を切り出すコードです。demo2フォルダー内にある、Youtubeなどからダウンロードして来た動画(india.mp4) から、静止画を切り出して image_dir フォルダーに格納します。

interval = 6 は静止画を切り出す間隔です。通常動画は 30フレーム/秒 なので、 interval = 6 だと5フレーム/秒(0.2秒に1フレーム)の静止画を切り出すことになります。

length = 60 は先頭から何枚静止画を切り出すかの指定です。

 

  

先程のコードを mov_png.py という名前で保存し、静止画を切り出したい動画( india.mp4 )を demo2 フォルダーに格納します。

detect_img フォルダーと image_dir フォルダーの中身は一端クリアしておいて下さい。

 

image_dir の中身

mov_png.pyを実行すると、image_dir 内に動画の静止画が格納されます。

 

  

detect_img フォルダーの中身

そして、ssd_model.py を実行すると、detect_img フォルダーの中に、物体検知をした画像が格納されます。

 

静止画からGIF動画を作成するコードです。detect_img フォルダーに格納されている静止画からGIF動画(india_ssd.gif) を作成します。このコードを demo2 フォルダーにmake_gif.py で保存します。

duration = 200 はフレームの再生速度で単位はミリ秒。200だと1フレームの再生が0.2秒となります。loop = 0 は再生をループさせる設定です。

 

india_ssd.gif

make_gif.py を実行すると、 demo2 フォルダーの中に、GIF動画(india_ssd.gif)が作成されます。

実は、静止画を物体検出しているだけですが、リアルタイム物体検出をしている様な雰囲気が出ますよね(笑)。

では、また。