🦮

YOLOv5-Nano と labelImg で犬猫検出モデル作ってみた

2024/11/26に公開

こんにちは!『麦ちゃーはん』です。

今回は、軽量で扱いやすい物体検出モデル「YOLOv5-Nano」と、シンプルなアノテーションソフト「labelImg」を使用して、犬と猫を検出できるモデルを作成しました。

簡易的な物体検出モデルを作成しただけですが、YOLOの学習の流れがつかめる内容になっていると思います。ぜひ最後までご覧ください。

labelImgについて

「labelImg」については、以前の記事で解説していますので、そちらをご覧ください。

https://zenn.dev/mugi_cha/articles/40ae064e1d2ceb

用意した画像

画像はこちらのものを利用させていただきました。
https://www.kaggle.com/datasets/tongpython/cat-and-dog?datasetId=23777&sortBy=voteCount

犬の画像

  • 学習用画像: 80枚
  • 検証用画像: 20枚

猫の画像

  • 学習用画像: 80枚
  • 検証用画像: 20枚

ディレクトリ構造

用意した画像と、labelImgで作成したアノテーションファイルのディレクトリ構造を紹介します。

dog_cat_dataset/
├── images/
│   ├── train/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   └──  val/
│      ├── image3.jpg
│       ├── image4.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── image1.txt
│   │   ├── image2.txt
│   │   └── ...
│   └── val/
│       ├── image3.txt
│       ├── image4.txt
│       └── ...
└── data.yaml  # データセットの設定ファイル

datasetsファルダの中にimagesファルダとlabelsファルダを作成します。

imagesファルダ内には学習用画像と検証用画像を入れるフォルダをそれぞれ作成し、labelsファルダ内には学習用画像と検証用画像のアノテーションファイルを入れるフォルダをそれぞれ作成します。

YOLOv5の導入方法

以下のコマンドで、YOLOv5の公式リポジトリをライブラリとして利用できるようにします。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 
pip install -U ultralytics

学習手順

設定ファイルの準備

以下の内容で、data.yamlを用意します。

# データセットのルートディレクトリ(全体の基準となるパス)
path: datasets/dog_cat_dataset  

# トレーニングデータの画像が保存されているサブディレクトリ
train: images/train  

# 検証データの画像が保存されているサブディレクトリ
val: images/val  

# クラス数(検出するオブジェクトの種類の数)
nc: 2  

# クラス名のリスト(検出対象のオブジェクト名を定義)
names: ['dog', 'cat']  

学習開始コマンド

yolo task=detect mode=train model=yolov5n.pt data=<data.yamlのパス> epochs=50 imgsz=640

各パラメータの説明

  • task: タスクタイプを指定する。detectは物体検出タスクを意味する。
  • mode: モードを指定。学習、検証、予測などのモードを選ぶことができる。
  • model: 学習に使用するYOLOのモデルを指定する。今回はYOLOv5-Nano。
  • data: データセットの設定ファイル(YAML形式)のパスを指定する。
  • epochs: 学習を何エポック(何周)実行するか指定する。
  • imgsz: 入力画像のサイズを指定する。

検証結果

モデル学習終了後に出力された検証画像は以下の通りです。

数か所で猫の誤検出が見られますが、現時点の学習量を考慮すれば、十分な精度と言えるでしょう。

検出精度を向上させる方法としては、学習データ数の増加やエポック数の増加、さらにはモデルの変更など、さまざまなアプローチが考えられます。しかし、今回はモデルを構築することが主な目的であるため、これをもって終了とさせていただきます。

おわりに

今回は、「YOLOv5-Nanoを用いた犬猫検出モデル」についてご紹介しました。

今後も引き続きYOLOを活用した取り組みを進めていく予定ですので、ぜひご期待ください。

最後までお読みいただき、心より感謝申し上げます!

Discussion