キカベン
機械学習でより便利な世の中へ
G検定対策
お問い合わせ
   

YOLOv5で転移学習をやってみる

thumb image

以前に、「ResNetで転移学習の方法を試してみる」や「転移学習の注意点」などで紹介した転移学習をYOLOv5でもやってみましょう。

YOLOv5は「YOLO v5で物体検出 (PyTorch Hubからダウンロード)」と同様にUltralyticsのものを使用します。

グーグルのOpen Imagesデータセットを使ってYOLOv5をもとに犬と猫の物体検出をするモデルの訓練をします。

データセットをYOLOv5での訓練用に準備できれば案外簡単に転移学習ができます。

1. VirtualenvでPython環境を作る🔝

今回は、VirtualenvでPythonの環境を作りましょう。

# プロジェクトのフォルダを作成し移動
mkdir yolov5-transfer-learning
cd yolov5-transfer-learning

# VirtualenvでPython環境を作りアクティベートする
python3 -m venv venv
source venv/bin/activate

# 一応、pipをアップデートしておく
pip install --upgrade pip

Virtualenvの良いところは、プロジェクトの中に環境を設定するフォルダを作成するので、一つのプロジェクトだけで使うときに分かりやすいところですね。

逆に、Condaで作る環境は複数プロジェクトで使うのに便利です。

2. Open Imagesからダウンロード🔝

まず、Open Imagesのデータセットをダウンロードするためのライブラリをインストールします。

pip install openimages

次に、犬と猫のデータセットを500画像ずつ(合計100画像)ダウンロードします。

oi_download_dataset --base_dir download --csv_dir download --labels Cat Dog --format darknet --limit 500

YOLOv5で使いやすいように、フォーマットでダークネットを指定します(–format darknet)。

また、これによって、後述するdarknet_obj_names.txtというファイルが作られます。

上記のコマンドを実行しても数分は凍りついたようになりますが、しばらくするとダウンロードが始まります。

やがてdownloadフォルダ内に犬と猫の画像と正解データが格納されます。

download folder

上図では、見やすくするために細かいファイルは省略していますが、実際にはdownloadフォルダにはdarknet_obj_names.txtというファイルがあって、中身はこんな感じです。

cat
dog

よって、インデックスで言うと、0が猫、1が犬になります。

上記のフォルダの構成のままではYOLOv5の訓練では使えないので、訓練用にデータセットを格納するためのフォルダを作りましょう。

3. YOLOv5の訓練用のフォルダ作成🔝

YOLOv5では下図のようなフォルダ構成を使います。

folder structure

方法はいろいろありますが、今回はPythonで次のようにフォルダを作ります。

import os

if not os.path.exists('data'):
    for folder in ['images', 'labels']:
        for split in ['train', 'val', 'test']:
            os.makedirs(f'data/{folder}/{split}')

上記のフォルダへ、画像と正解データを振り分けていけばYOLOv5の訓練で使うことができます。

4. 画像ファイル名の重複チェック🔝

念のためファイル名に重複がないかチェックしておきましょう。

フォルダを指定してファイル名をset入れて取り出します。

import glob

def get_filenames(folder):
    filenames = set()
    
    for path in glob.glob(os.path.join(folder, '*.jpg')):
        # pathを分解してファイル名だけを取り出す
        filename = os.path.split(path)[-1]        
        filenames.add(filename)
        
    return filenames

# 犬と猫の画像ファイル名のセット
dog_images = get_filenames('download/dog/images')
cat_images = get_filenames('download/cat/images')

ファイル名の重複をチェックします。

# 同じ名前のファイル名をチェック
duplicates = dog_images & cat_images

print(duplicates)

ここで、{'1417eccd5854e04a.jpg', '0dcd8cc4b35a93b4.jpg', '0838125199f2caa7.jpg'}がプリントされました。

3つのファイル名が犬と猫のフォルダで重複しています。

数が少ないでの目視で画像をチェックしてみます。

from PIL import Image

# 同じファイル名の画像をcatとdogのフォルダから表示してみる
for file in duplicates:
    for animal in ['cat', 'dog']:
        Image.open(f'download/{animal}/images/{file}').show()
yolov5 cat
yolov5 cat
yolov5 cat

同じ猫の画像がdogフォルダに紛れ込んでいたようです。

理由は分かりませんが、追求しても仕方ないので取り除きましょう。

dog_images -= duplicates

print(len(dog_images))

497とプリントされて3つのファイル名が犬画像ファイル名のセットから取り除かれたのが確認できました。

5. データセットを振り分ける🔝

犬と猫のデータを訓練用、バリデーション、テスト用に振り分けます。

画像データは特に規則性もなく、ファイルのリストをシャッフルする必要はないかもしれませんが、一応やっておきます。

import numpy as np

dog_images = np.array(list(dog_images))
cat_images = np.array(list(cat_images))

# 乱数シードを固定して再現性を確保
np.random.seed(42)
np.random.shuffle(dog_images)
np.random.shuffle(cat_images)

下記のコードはちょっと長いですが、画像と正解データのファイルを振り分けてコピーしているだけです。

import shutil

def split_dataset(animal, image_names, train_size, val_size):
    for i, image_name in enumerate(image_names):
        # 正解データのファイル名
        label_name = image_name.replace('.jpg', '.txt')
        
        # train, val, testを区分けする
        if i < train_size:
            split = 'train'
        elif i < train_size + val_size:
            split = 'val'
        else:
            split = 'test'
        
        # データのコピー元
        source_image_path = f'download/{animal}/images/{image_name}'
        source_label_path = f'download/{animal}/darknet/{label_name}'

        # データのコピー先
        target_image_folder = f'data/images/{split}'
        target_label_folder = f'data/labels/{split}'

        # コピーする
        shutil.copy(source_image_path, target_image_folder)
        shutil.copy(source_label_path, target_label_folder)

# 猫
split_dataset('cat', cat_images, train_size=400, val_size=50)

# 犬は3つ足りないので1つずつ減らす
split_dataset('dog', dog_images, train_size=399, val_size=49) 

これでデータの振り分けは出来たのですが、次に正解データの中身を理解しておきましょう。

6. 正解データの中身🔝

ちなみに、正解データ(*.txt)の中身はこんな感じです。

1 0.36750000000000005 0.6512604999999999 0.46125000000000005 0.402427
1 0.806875 0.31792699999999996 0.24625000000000008 0.294118

最初は0か1で、猫か犬かを表しています。

次の4つの数値はでBounding Boxの中央のx値とy値、幅wと高さhを画像のサイズに対して相対的に(0~1)で表しています。

試しに訓練用のデータで表示してみましょう。

from PIL import Image, ImageDraw

def show_bbox(image_path):
    # image_pathのフォルダ名と拡張子を変更してラベルファイルのパスを作る
    label_path = image_path.replace('/images/', '/labels/').replace('.jpg', '.txt')

    # 画像を開き、描画ようにImageDrawを作る
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    with open(label_path, 'r') as f:
        for line in f.readlines():
            # 一行ごとに処理する
            label, x, y, w, h = line.split(' ')

            # 文字から数値に変換
            x = float(x)
            y = float(y)
            w = float(w)
            h = float(h)

            # 中央位置と幅と高さ => 左上、右下位置
            W, H = image.size
            x1 = (x - w/2) * W
            y1 = (y - h/2) * H
            x2 = (x + w/2) * W
            y2 = (y + h/2) * H

            # BoundingBoxを赤線で囲む
            draw.rectangle((x1, y1, x2, y2), outline=(255, 0, 0), width=5)

    image.show()
    

show_bbox('data/images/train/00a1ab47b5439e8c.jpg')

こんな感じになります。

7. 転移学習の準備🔝

データセットの準備ができたので、転移学習の準備に入ります。

まず、YOLOv5のレポジトリをクローンしてから、必要なライブラリをインストールします。

yolov5-transfer-learningのフォルダの中で以下を実行してください。

git clone https://github.com/ultralytics/yolov5

pip install -U -r yolov5/requirements.txt

今回も小さいモデルであるyolov5sを使います。

8. バックボーンを固定する🔝

次に、YOLOv5のモデルのウェイトを固定して転移学習中に変更されないようにします。

このセクションは、Transfer Learning with Frozen Layersを参考にしています。

yolov5/models/yolov5s.yamlを見ると最初の10個の層が特徴量を取り出すバックボーン(Backbone)になっています。

# YOLOv5 backbone 
 backbone: 
   # [from, number, module, args] 
   [[-1, 1, Focus, [64, 3]],  # 0-P1/2 
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4 
    [-1, 3, BottleneckCSP, [128]], 
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8 
    [-1, 9, BottleneckCSP, [256]], 
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16 
    [-1, 9, BottleneckCSP, [512]], 
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32 
    [-1, 1, SPP, [1024, [5, 9, 13]]], 
    [-1, 3, BottleneckCSP, [1024, False]],  # 9 
   ] 
  
 # YOLOv5 head 
 head: 
   [[-1, 1, Conv, [512, 1, 1]], 
    [-1, 1, nn.Upsample, [None, 2, 'nearest']], 
    [[-1, 6], 1, Concat, [1]],  # cat backbone P4 
    [-1, 3, BottleneckCSP, [512, False]],  # 13 
  
    [-1, 1, Conv, [256, 1, 1]], 
    [-1, 1, nn.Upsample, [None, 2, 'nearest']], 
    [[-1, 4], 1, Concat, [1]],  # cat backbone P3 
    [-1, 3, BottleneckCSP, [256, False]],  # 17 (P3/8-small) 
  
    [-1, 1, Conv, [256, 3, 2]], 
    [[-1, 14], 1, Concat, [1]],  # cat head P4 
    [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium) 
  
    [-1, 1, Conv, [512, 3, 2]], 
    [[-1, 10], 1, Concat, [1]],  # cat head P5 
    [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large) 
  
    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5) 
   ] 

よって、この10層を固定するために、yolov5/train.pyを開いて以下の部分を見つけます。

# Freeze
freeze = []  # parameter names to freeze (full or partial)

これを次のように変更してください。

freeze = ['model.%s.' % x for x in range(10)]

これだけで、最初の10層のウェイトが固定され訓練中に変更されなくなります。

9. 設定YAMLファイルを作る🔝

訓練する際に必要な設定YAMLファイルを作ります。

yolov5-transfer-learningのフォルダの中にcats_and_dogs.yamlというファイルを作成します。

# データセットの場所
train: data/images/train
val:   data/images/val
test:  data/images/test

# クラスの数(nc = number of classes)
nc: 2

# クラスの名前(0が猫、1が犬)
names: ['cat', 'dog']

これで訓練の準備ができました。

10. 訓練を実行する🔝

以下のようにして訓練を実行します。

python yolov5/train.py --data cats_and_dogs.yaml --weights yolov5s.pt --epochs 100 --batch 4

バッチサイズは使用するマシンのスペックに合わせて調節してください。

エポックは100も必要ないかもしれませんが、そのくらいあれば十分でしょう。

11. Tensorboardで状況を確かめる🔝

Tensorboardを使って訓練の状況を見ることもできます。

tensorboard --logdir runs/

ブラウザーからlocalhost:6006を開けばTensorboardのチャートなどを閲覧することができます。

12. 訓練後のモデル評価🔝

訓練が終わると、runs/train/exp/weightsにモデルのウェイトのファイルが二つ格納されています。

  • best.pt 一番成績が良かったモデル
  • last.pt 一番最後のエポックのモデル

ちなみにruns/train/exp/weightsexpexperiment(実験)の略で、訓練を実行するたびに新しいフォルダ (exp2exp3、..)と自動で作られます。

では、これらのモデルのウェイトを使ってテストのデータセットで評価をしましょう。

python yolov5/test.py --data cats_and_dogs.yaml --weights runs/train/exp/weights/best.pt

以下のような結果が出ました、なかなか良いですね。

Class Images Labels P     R     mAP@.5 mAP@.5:.95: 
  all     99    113 0.743 0.81  0.827  0.56
  cat     99     51 0.754 0.781 0.808  0.492
  dog     99     62 0.732 0.839 0.846  0.628

参照【物体検出】mAP (mean Average Precision)とは? その目的と計算方法を解説します

13. 事前処理のコードのまとめ🔝

import glob
import shutil
import os
import numpy as np
from PIL import Image, ImageDraw


# データセットを振り分けるためのフォルダを作る
if not os.path.exists('data'):
    for folder in ['images', 'labels']:
        for split in ['train', 'val', 'test']:
            os.makedirs(f'data/{folder}/{split}')

            
def get_filenames(folder):
    filenames = set()
    
    for path in glob.glob(os.path.join(folder, '*.jpg')):
        # pathを分解してファイル名だけを取り出す
        filename = os.path.split(path)[-1]        
        filenames.add(filename)
        
    return filenames


# 犬と猫の画像ファイル名のセット
dog_images = get_filenames('download/dog/images')
cat_images = get_filenames('download/cat/images')


# 同じ名前のファイル名をチェック
duplicates = dog_images & cat_images
print(duplicates)


# 重複画像をマニュアルでチェック
for file in duplicates:
    for animal in ['cat', 'dog']:
        Image.open(f'download/{animal}/images/{file}').show()


# 重複画像ファイル名を取り除く
dog_images -= duplicates

print(len(dog_images))


# セット(set)からNumpy配列に変換
cat_images = np.array(list(cat_images))
dog_images = np.array(list(dog_images))


# 乱数シードを固定して再現性を確保し、シャッフルする
np.random.seed(42)
np.random.shuffle(cat_images)
np.random.shuffle(dog_images)


def split_dataset(animal, image_names, train_size, val_size):
    for i, image_name in enumerate(image_names):
        # 正解データのファイル名
        label_name = image_name.replace('.jpg', '.txt')
        
        # train, val, testを区分けする
        if i < train_size:
            split = 'train'
        elif i < train_size + val_size:
            split = 'val'
        else:
            split = 'test'
        
        
        # データのコピー元
        source_image_path = f'download/{animal}/images/{image_name}'
        source_label_path = f'download/{animal}/darknet/{label_name}'

        # データのコピー先
        target_image_folder = f'data/images/{split}'
        target_label_folder = f'data/labels/{split}'

        # コピーする
        shutil.copy(source_image_path, target_image_folder)
        shutil.copy(source_label_path, target_label_folder)


split_dataset('cat', cat_images, train_size=400, val_size=50)
split_dataset('dog', dog_images, train_size=399, val_size=49) # 3つ足りないので1つずつ減らす


# BoundingBoxを表示してみる
def show_bbox(image_path):
    label_path = image_path.replace('/images/', '/labels/').replace('.jpg', '.txt')

    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    with open(label_path, 'r') as f:
        for line in f.readlines():
            # 一行ごとに処理する
            label, x, y, w, h = line.split(' ')

            # 文字から数値に変換
            x = float(x)
            y = float(y)
            w = float(w)
            h = float(h)

       # 中央位置と幅と高さ => 左上、右下位置
            W, H = image.size
            x1 = (x - w/2) * W
            y1 = (y - h/2) * H
            x2 = (x + w/2) * W
            y2 = (y + h/2) * H

            # BoundingBoxを赤線で囲む
            draw.rectangle((x1, y1, x2, y2), outline=(255, 0, 0), width=5)

    image.show()
    

show_bbox('data/images/train/00a1ab47b5439e8c.jpg')

今日はここまで。



コメントを残す

メールアドレスは公開されません。