YOLOv8でオリジナルデータを学習させるまで

YOLOv8で自身が用意した画像を使って、学習させるまでにはいくつかのステップがある。

  1. アノテーション(ラベル付)
  2. 学習させるためのフォルダの作成と画像の配置
  3. yamlファイルの作成
  4. 学習の実行

アノテーション

YOLOv8では、画像から推論を行う場合には、その画像中のどこにどんなものがあるかを教える必要がある。そのために行うのがアノテーションである。画像中の”場所”と”物”を対応させて保存することで、それを学習で利用できるようにする。今回アノテーションでは無料でも使用できるFastLabel(https://fastlabel.ai/)を使用した。抵抗の推定では画像中から抵抗の場所を特定できればいいので、アノテーションクラスを1つだけ用意し、アノテーションを行なった。

FastLabelの使い方などは公式で出しているので、ここでは省略する。
公式ドキュメント: https://fastlabel.notion.site/FastLabel-Docs-8eac0c2dc57540ee87eedd409fad0e15
動画: https://www.youtube.com/watch?v=TWtXZe5VJkQ

アノテーションが完了したら、エクスポート画面で出力形式をYOLOに設定し画像出力をオリジナルに設定した後、エクスポートを実行する。その後、エクスポートされたzipファイルをダウンロードする。ダウンロードしたzipファイルを展開すると、以下のような構造になっている。

展開フォルダ
|- originals
| |- 画像ファイル
|
|- yolo
|- class.txt
|- annotations
|- txtファイル

画像ファイルとtxtファイルは対応した名前になっており、これらを使って学習を行う

画像の配置とyamlファイルの作成

アノテーションを行った画像とtxtファイルをGoogle driveにアップロードする。今回はアップロードにあたって、yolov8フォルダの下にdataフォルダを作成し、その下に画像とラベルが保存されているフォルダごとアップロードした。以下に構造を示す。

yolov8
|- data
|- originals
| |- 画像ファイル
|
|- yolo
|- class.txt
|- annotations
|- txtファイル

アップロードしたファイルから学習用データ(train)と検証用データ(val)を分けて保存する。その時のフォルダ構造は以下のようにする。

yolov8
|- data
|- resistance
|- labels
| |- train
| | |- txtファイル
| |- val
| |- txtファイル
|
|- images
|- train
| |- 画像ファイル
|- val
|- 画像ファイル

学習用データと検証用データをそれぞれ保存したあと、yamlファイルにフォルダ構造やクラス数などを記述する。yamlファイルの書き方はyolov8に例がある(https://docs.ultralytics.com/ja/datasets/detect/)。
学習用と検証用にデータを分け保存し、yamlファイルに書き込むのは大変なので、これらの処理を行うプログラムを作成した。このプログラムを実行することで、yamlファイルまでの作成を行う(yamlファイルはyolov8/data/ に配置される)。

Python
import glob
import os
import shutil
import random

# パラメータ
CURRENT_FOLDER = '/content/drive/MyDrive/yolov8/data/' # 画像などを配置するフォルダ
VAL_IMG_RATE = 20.0 # 検証用画像の割合 [%]
IMG_EXT = '.JPG' # 画像ファイルの拡張子
CLASS_NUM = 1 # 分類数
YAML_FILE_NAME = 'resistance' # yamlファイルの名前


# 学習用画像(train)と検証用画像(val)を分けて保存
data_set_folder = CURRENT_FOLDER+YAML_FILE_NAME+'/'
# 既存フォルダの削除
if os.path.exists(data_set_folder+'images'):
  shutil.rmtree(data_set_folder+'images')
if os.path.exists(data_set_folder+'lables'):
  shutil.rmtree(data_set_folder+'labels')
# フォルダ作成
os.makedirs(data_set_folder+'images/train',exist_ok=True)
os.makedirs(data_set_folder+'images/val',exist_ok=True)
os.makedirs(data_set_folder+'labels/train',exist_ok=True)
os.makedirs(data_set_folder+'labels/val',exist_ok=True)

# 画像ファイルのファイル名取得
img_names = []
for filename in os.listdir(CURRENT_FOLDER+'originals/'):
  # ファイルが画像ファイルであることを確認
  if filename.endswith(IMG_EXT):
    # 拡張子を除いてファイル名を取得
    img_name, _ = os.path.splitext(filename)
    img_names.append(img_name)

# シャッフル
img_names = random.sample(img_names, len(img_names))

#学習用データと検証用データに分離
split_index = int(len(img_names) * VAL_IMG_RATE/100)
val_img_names = img_names[:split_index]
train_img_names = img_names[split_index:]

# データのコピー
for val_img_name in val_img_names:
  shutil.copy(CURRENT_FOLDER+'originals/'+val_img_name+IMG_EXT, data_set_folder+'images/val/')
  shutil.copy(CURRENT_FOLDER+'yolo/annotations/'+val_img_name+'.txt', data_set_folder+'labels/val/')

for train_img_name in train_img_names:
  shutil.copy(CURRENT_FOLDER+'originals/'+train_img_name+IMG_EXT, data_set_folder+'images/train/')
  shutil.copy(CURRENT_FOLDER+'yolo/annotations/'+train_img_name+'.txt', data_set_folder+'labels/train/')

#yamlファイルの作成
with open(CURRENT_FOLDER+YAML_FILE_NAME+'.yaml', 'w') as file_data:
  file_data.write('path: '+CURRENT_FOLDER+YAML_FILE_NAME+'\n')
  file_data.write('train: images/train \n')
  file_data.write('val: images/val \n')
  # file_data.write('nc: '+str(CLASS_NUM)+' \n')
  file_data.write('names: \n')
  for i in range(CLASS_NUM):
      file_data.write('   '+str(i)+': \n')

学習の実行

学習のさせ方は公式にも載っているので詳しく知りたい方はそっちを見てほしい(https://docs.ultralytics.com/ja/modes/train/)。
ここではGoogle Colaboratoryを使った学習のさせ方と例を説明する。

まずGoogle ColaboratoryのランタイムのタイプをCPUからGPUに変更する。CPUでも学習はできるが、GPUを使用して学習させた方が、圧倒的に早く学習が終わる。また、Google Colaboratoryでは無料でGPUが利用できるので、GPUを使用した方がいい。GPUにするには右上の下三角のボタンをクリックし、ランタイムのタイプを変更をクリックすることで、ランタイムのタイプを変更の画面が出てくるので、そこでGPUを選択する。

GPUに変更した後は、学習のプログラムを実行する。dataでは作成したyamlファイルのパスを指定する。今回は以下のような設定で学習を行なった。

Python
%cd /content/drive/MyDrive/yolov8/
from ultralytics import YOLO
model = YOLO("yolov8s.yaml")
results = model.train(data="/content/drive/MyDrive/yolov8/data/resistance.yaml", epochs=300, imgsz=640)

学習の結果

学習が完了するとyolov8/runs/detect/train/に結果が保存されている。その中のresults.pngを見ると学習によって誤差(loss)が低下しているのがわかる。lossが低下せずにprecisionが上昇していない場合には、学習が失敗している可能性があるため注意が必要である。学習後の推定にはweightsフォルダ下の.ptを利用するので消去しないように注意する。

コメントする