人工知能の技術

PyTorchで”ディープ”な物体検出:学習済みモデルを実行する(画像)

こんにちは。現役エンジニアの三年坊主(@SannenBouzu)です。

エンジニア歴は6年。うち2年以上のAI開発を経験し、Web技術とAI技術をバランスよく習得してきました。

今回は、PyTorchの学習済みモデルを実行し、画像から物体検出を行う方法を紹介します。

PyTorchの基礎、物体検出の基礎については、別記事で取り扱う予定です

PyTorchの基礎①: ニューラルネットワークをTensorで実装【PyTorchの基礎】簡単なニューラルネットワークをTensorを使って実装します。torch.nnを使わず、処理の流れを具体的に追っていくので理解が深まります。その知識は、自然言語処理、画像処理など、自分の用途に合わせて実装する際に役立つでしょう。...

「学習済みモデル」を使うことで、「データ収集」「前処理」「学習」といった、ディープラーニングで時間がかかる多くのプロセスをスキップして、「モデルを使った推論」という結果だけを素早く得ることができます。

また、記事で扱う「モデルを定義→入力を与える→出力を得る」という「推論」の流れは、物体検出に限りません。一度理解しておくと、ディープラーニング、もっと広く機械学習全般にも応用が可能です。

記事のまとめ部分では、全体像を整理しやすいように、ObjectDetectorというクラスを定義していますので、復習に役立ててください。

 

この記事はこんな方におすすめ!

  • PyTorchを使った物体検出をサクッと試してみたい方
  • Pythonの物体検出プログラムを確認したい方
Pythonを快適に使いこなすMac環境【現役エンジニアおすすめはPro 13インチ】MacでPythonプログラミングをしたい人向けに、現役エンジニアおすすめMac環境を紹介します。2011年からMacBook Proを使い続け、Mac買い替えやPythonインストールを経験している私が実体験に基づいて説明します。ぜひご覧ください。...

ライブラリのインポート(PyTorchや画像処理ライブラリなど)

環境の準備については、別記事で取り扱う予定です。Google Colaboratory、またはDockerでの環境構築をおすすめします

PyTorchなど、今回使うライブラリをインポートします。

  • PyTorch:Pythonのオープンソースの機械学習ライブラリ
  • torchvision:PyTorchのパッケージ。コンピュータビジョンにおける有名なデータセット・モデルアーキテクチャ・画像変換処理などから構成される
  • Pillow(PIL):Pythonの画像処理ライブラリ
  • OpenCV(cv2):Pythonの画像処理ライブラリ

 

今回の記事は、こちらのバージョンで動作確認しています。

 

GPU/CPUの確認と切り替え

PyTorchといえばディープラーニング。ディープラーニングといえば、主に画像処理を高速化するためにGPU (Graphics Processing Unit) を使う場面が増えています。

GPUを使えるかどうか確認しておきましょう。ちなみに、今回はGPUを使えなくても問題ありません(私も手元のMacで試しました)。

 

torch.deviceを使って、PyTorchのテンソルをどのデバイスに割り当てるのか指定します(CPUかGPUか、GPUが複数あるときにはどのGPUか)。

 

PyTorchのインポートを確認する

順番が入れ変わってしまいましたが、適当なテンソルを作って、PyTorchのインポートができているか確認しましょう。

 

物体検出に使う画像とモデルを用意する

物体検出をするため、適当な画像とモデルを用意します。

画像

飛行機と時計の画像を使ってみます。

学習済みのモデル

torchvisionは、コンピュータビジョンのタスクに便利なデータセット、モデル、画像変換処理などをまとめたパッケージです。

torchvisionには、Faster R-CNNという物体検出モデルが用意されているので、使ってみましょう。

物体検出の基礎については、別記事で取り扱う予定です

pretrained=Trueにして、学習済みのモデルを利用します。

 

モデルをdevice(CPUまたはGPU)に割り当てます。

 

PIL画像をテンソルに変換する

PILライブラリで読み込んだ画像を、PyTorchの物体検出モデルで使えるよう、テンソルに変換する必要があります。

 

img1, img2のsizeを確認しましょう。

 

物体検出モデルを実行する(予測・推論)

公式ドキュメントで、Faster R-CNNの入力を確認します。

The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each image, and should be in 0-1 range. Different images can have different sizes.

先ほどテンソルに変換したimg1とimg2のサイズは、

  • (3, 1982, 3027)
  • (3, 2560, 3840)

どちらも、C(チャネル), H(高さ), W(幅)の順になっていましたね。

値の範囲も、0から1の間に収まっています。

 

モデルに入力する画像(テンソル)には問題なさそうなので、さっそくモデルに入力してみましょう。

  • model.eval():モデルをevaluationモードに切り替える
    • →学習時特有の処理を止める
  • torch.no_grad():自動勾配の計算を止める
    • →メモリ使用量を減らせる
  • 画像のテンソルは、modelの引数にリストで与える
    • →predictionsは画像ごとの物体検出結果を含むリスト

 

 

物体検出モデルの予測結果には、「スコア」が含まれています。confidenceという呼び方もありますが、要するに「どの程度自信を持てる予測結果か」という数字を、0から1の間で出してくれます。

スコアが低い予測結果を受け入れる場合、より多くの物体を検出できる可能性が上がる一方で、正しくない物体検出が増えてしまうかもしれません

上記のコードでは、score_thresholdに設定した値より低いスコアの予測結果をはじいています。

 

predictionsに格納した、予測結果を見てみましょう。

 

公式ドキュメントで、Faster R-CNNの出力を確認します。

  • boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values of x between 0 and W and values of y between 0 and H
  • labels (Int64Tensor[N]): the predicted labels for each image
  • scores (Tensor[N]): the scores or each prediction

boxesというのが、バウンディングボックス(bounding box)、物体をちょうど囲むのに必要な大きさの長方形(矩形)のことで、モデルが予測する物体の位置を表しています。

画像を入力すると、画像のどの位置に、何が写っているのか予測する「物体検出」を実行することができました。

 

バウンディングボックス(物体の位置)を可視化する

物体検出ができたので、人間の目にも分かりやすい形で可視化してみましょう。

さて、今回使った学習済みモデルは、COCO Datasetというデータセットを使って「事前に学習した物体」を検出できるものです。

「事前に学習した物体」というのは、例えば

  • 人物
  • 自転車
  • オートバイ
  • 飛行機

など。具体的には以下の通りで、

N/Aと__background__を除くと80種類になっています。

 

predictionsのlabelsに

  • ‘labels’: tensor([5])
  • ‘labels’: tensor([85])

と出ていたのは、画像から予想できる通り、それぞれ「飛行機」と「時計」でした。

 

cv2.rectangle

続いて、物体の位置を可視化していきましょう。

いくつか方法はありますが、こちらはOpenCVのrectangle関数を使って長方形を描画する方法です。

  • OpenCVの関数を使うので、画像もOpenCVで読み込んでいる
  • cv2.rectangleは座標をintで指定しないといけないので、predictionsのboxesをintに変換している
  • cv2.putTextでラベルの名前(airplane, clockなど)を描画

 

まとめ: ObjectDetectorクラスを定義して全体像を整理する 

今回は、PyTorchの学習済みモデルを実行し、画像から物体検出を行う方法を紹介しました。

最後に、今回の記事で紹介した内容をObjectDetectorというクラスにまとめます。

 

使い方は以下の通り。

「モデルを定義→入力を与える→出力を得る」という「推論」の流れを、実感していただけたかなと思います。

 

ところで、学習済みモデルを使うと、手軽に物体検出をできる反面、当然ながら、「事前に学習した物体しか検出できない」というデメリットがあります。

自分で用意したデータ・誰かが用意してくれたデータでモデルを学習させ、好きな物体を検出する方法については、別記事で取り扱う予定です

▼経験の棚卸しで納得のキャリアを▼
▼自宅で簡単・お得にふるさと納税▼
【期間限定】無料登録でプレゼント

COMMENT

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください