Pythonによるデータ分析・機械学習ブログ

Pythonによるデータの前処理・グラフ化、機械学習、統計解析、画像処理、Webスクレイピング、自然言語処理の基礎について紹介していきます。

MENU

【第19回】YOLOv9で自作データセットを用いて学習

はじめに

 近年、AI(Artificial Intelligence:人工知能)の急速な発展に伴い、それらを用いた技術が活用され始めています。例えば、AIによる画像認識技術として物体検出などがあります。物体検出は「画像の中から特定の物体の位置、種類、個数などの情報を認識する技術」であり、AIを用いて比較的高速に画像内の物体情報を取得することができます。この技術を用いて製造業では品質評価、外観検査、異常検知などに活用されています。

 リアルタイムに物体検出を行うアルゴリズムにYOLO(You Only Look Onse)と呼ばれるものがあり、CNN(Convolutional Neural Network:畳み込みニューラルネットワーク)を用いたアルゴリズムとなっています。2024年2月に最新の物体検出モデルYOLOv9が公開されています。

 本記事では、YOLOv9で自作データセットを用いて学習を行う方法を簡単に紹介します。

推論デモについての内容は下記の記事をご覧ください。

chantastu.hatenablog.com

 

1. YOLOv9のリポジトリ

YOLOv9のGithubリポジトリは以下です。

github.com

 

2. 動作環境構築

以下のコマンドで動作環境を構築できます。

 

3. データセットの用意

Roboflowで学習のテストを実施するためのサンプルデータセットを用意しました。アノテーション内容としては画像20枚に対してpersonとpotted plantの2クラスに対して矩形付与をしており、学習:検証:テスト=7:2:1の割合になっています。下記がアノテーション画像例です。

 

下記リンクからZipファイルがダウンロードできます。

https://app.roboflow.com/ds/Cmu0SZTXY7?key=VwBP4N8NgZ

また、以下のコマンドでディレクトリにダウンロードできます。

次にダウンロードした中のdata.yamlファイルの一部を修正します。

data.yamlを開くと、下記のようにtrain、valid、testのパスが書かれた行があるかと思います。

この部分のパスを修正します。上記でデータセットを「datasets/detect-test」内にダウンロードしたので、その部分を反映させます。以下のように内容を修正します。

 

4. 事前学習済みモデルの取得

 次に推論で使用するYOLOv9で提供されている事前学習済みモデル「gelan-c.pt」をダウンロードします。今回はyolov9ディレクトリ直下にweightsフォルダを作成し、そこにモデルをダウンロードします。

 

5. 学習実行

 それでは学習を実行していきます。下記のコマンドを実行することで学習が実行されます。--batchはバッチ数、--epochはエポック数を指定できるので適宜変更してください。

実行すると下記のようなログが表示されます。最後の方を見ると「runs/train/exp8/weights/best.pt」のパスに学習時の精度が良かったモデルが保存されていることが確認できます。

 

6. 検証

次にバリデーションをしていきます。ここでは学習したモデルの品質を評価します。下記のコマンドで実行します。--batchはバッチ数、--confは信頼度スコアの閾値、--iouはNMSの閾値を指定できるので適宜変更してください。

--dataは先ほど修正したdata.yaml相対パスを、--weightsは学習したモデルが格納された相対パスを指定します。

上記プログラムを実行すると下記内容が表示されます。

 

7. カスタムモデルで推論

最後にカスタムモデルで推論をしていきます。下記のコマンドで推論を実行します。ここで--confは信頼度スコア、--weightsには学習したモデルの相対パス、--sourceにはテスト画像が格納されたディレクト相対パスやテスト画像の相対パスを指定します。

推論実行後、「/runs/detect/exp{No.}※No.は任意の数字」に推論結果の画像が保存されます。今回の結果は下図のようになります。

 

オススメ書籍

物体・画像認識と時系列データ処理入門

現場で使える!PyTorch開発入門 深層学習モデルの作成とアプリケーションへの実装

PyTorchで作る!深層学習モデル・AIアプリ開発入門

 

おわりに

 本記事では、YOLOv9で自作データセットを用いて学習を行う方法を簡単に紹介しました。最新のYOLOv9は非常に使いやすい印象です。アノテーション・学習をご自身でも行ってみて、目的達成のために物体検出を活用してみてください。