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

この記事で学ぶこと

この記事では、YOLOv5で転移学習を行う手順を解説します。以前に、「ResNetで転移学習の方法を試してみる」や「転移学習の注意点」などで紹介した転移学習をYOLOv5でもやってみましょう。

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

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

データセットをYOLOv5での訓練用に準備できれば案外簡単に転移学習ができます。では、さっそく始めましょう。

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 (Mini Conda)では複数プロジェクトで同じ環境を使うのに便利です。

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フォルダにはdarknet_obj_names.txtというファイルがあって、中身はこんな感じです。

cat
dog

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

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

YOLOv5の訓練用のフォルダ作成

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

方法はいろいろありますが、今回は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の訓練で使うことができます。

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

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

フォルダを指定してファイル名を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'}がプリントされました。

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

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

from PIL import Image

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

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

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

dog_images -= duplicates

print(len(dog_images))

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

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

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

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

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つ足りないので１つずつ減らす
split_dataset('dog', dog_images, train_size=399, val_size=49) 

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

正解データの中身

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

1 0.36750000000000005 0.6512604999999999 0.46125000000000005 0.402427
1 0.806875 0.31792699999999996 0.24625000000000008 0.294118

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

次の４つの数値はで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')

こんな感じになります。

転移学習の準備

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