【Python×物体検出】超簡単！YOLOv8による物体の学習から推論まで

google driveのディレクトリ構成

マイドライブに「data」フォルダを作成しておく

content
    |_____drive
            |_____My Drive
                    |_____data

作業ディレクトリの作成

PCで画像の収集やアノテーションを行うため、作業ディレクトリを作成しておく

python_yolov5
	|__________images
	|          	|__________colab_upload
	|          	|          	|__________test
	|          	|          	|__________train
	|          	|          	|__________val
	|          	|
	|          	|__________resource
	|
	|__________1_get_img.py
	|__________2_rename.py
	|__________3_DataAugumentation.py
	|__________4_shuffle.py

学習したい物体が映った画像を収集

1_get_img.pyの中身

from icrawler.builtin import GoogleImageCrawler
crawler = GoogleImageCrawler(storage={"root_dir": "resourceフォルダのパス"})

#keyword => 入手したい画像のキーワード
#max_num => 入手する画像の枚数
crawler.crawl(keyword='腕時計', max_num=100)
print('-----finifh-----')

resourceフォルダの中に画像が保存されているはずです。

自前の画像データがある場合は1_get_img.pyを実行せずにresourceフォルダに格納するだけで良いです。

画像のリネーム

2_rename.pyの中身

import os
import glob

path = "resourceフォルダのパス"
files = glob.glob(path + '/*')

#画像の名前を番号に変更し.pngで保存する
for i, f in enumerate(files, 1):
    os.rename(f, os.path.join(path, '{0:05d}'.format(i) + '.png'))
    #print(i)

画像を扱いやすくするために名前を番号に、拡張子をpngに変更しています。

学習に用いるデータは数が多いのでpythonで実行した方が良いです。

データの拡張

3_DataAugumentation.pyの中身

import cv2
import pathlib

input_dir = "resourceフォルダのパス"
output_dir = "resourceフォルダのパス"
image_list_png = list(pathlib.Path(input_dir).glob('**/*.png'))

rot_angle=180

for i in range(len(image_list_png)):
    img_png = cv2.imread(str(image_list_png[i]), cv2.IMREAD_COLOR)
    h,w=img_png.shape[:2]
    ROT= cv2.getRotationMatrix2D(center=(w/2,h/2),angle=rot_angle,scale=1)
    img_png = cv2.warpAffine(img_png,ROT,dsize=(w,h))
    output_path = output_dir + '/'  + 'add_ud_' + image_list_png[i].name 
    cv2.imwrite(output_path, img_png)
    #print('ud_' + str(i))

for i in range(len(image_list_png)):
    img_png_flip = cv2.imread(str(image_list_png[i]), cv2.IMREAD_COLOR)
    img_flip_lr = cv2.flip(img_png_flip, 1)
    output_path = output_dir + '/'  + 'add_fp_' + image_list_png[i].name
    cv2.imwrite(output_path, img_flip_lr)
    #print('fp_' + str(i))

print('----------finish----------')

より正確に物体の推論を行うためにはより多くのサンプルが必要になります。

しかし、大量のサンプルデータを収集するのは非常に大変な作業のため、用意した画像を上下反転、左右反転してデータを3倍に増やします。

3_DataAugumentation.py実行後、resourceフォルダの中身が増えているはずです。

アノテーションを行う

アノテーションツールを使用してアノテーションを行います。

今回は「labelImg」を使用します。

コマンドプロンプトにてインストールします。

pip3 install labelImg

インストールできたらコマンドプロンプトから実行します。

labelImg

「labelImg」を用いたアノテーションのやり方については省略します。

画像データが格納されているresourceファイルに各画像と同じ名前で保存していきます。

出力形式をYOLOに変更するのを忘れないでください。

アノテーション後のresourceフォルダの中身

resource
	|__________00001.png
	|__________00001.txt
	|__________00002.png
	|__________00002.txt
	|

画像の振り分け

これまでに用意した画像をtrain, test, valに分けて別々のフォルダに移動させていきます。

今回はtrain：test：val＝7：1：2という比率で分けていきます。

4_shuffle.pyの中身

import glob
import os
import random
import shutil
import numpy as np

input_dir = "resourceフォルダのパス"
save_dir = "colab_uploadフォルダのパス"

def split_data(folder,savefolder,train_rate,val_rate):
    imgs = glob.glob(os.path.join(folder,"*.png"))
    txts = glob.glob(os.path.join(folder,"*.txt"))

    imgs_index = np.array(range(len(imgs)))
    imgs_index_shuffle = np.random.permutation(imgs_index)
    n = len(imgs)
    train_n = int(n * train_rate)
    val_n = int(n * val_rate)
    test_n = n - train_n - val_n
    print(train_n,test_n,val_n)

    for i, index  in enumerate(imgs_index_shuffle):
        img_path = imgs[index]
        txt_path = txts[index]
        img_name = img_path.split("\\")[10]
        txt_name = txt_path.split("\\")[10]
        print(img_name,txt_name)
        if i < train_n:
            shutil.copyfile(img_path, savefolder + f"/train/{img_name}")
            shutil.copyfile(txt_path, savefolder + f"/train/{txt_name}")
        elif train_n <= i < train_n + val_n:
            shutil.copyfile(img_path, savefolder + f"/val/{img_name}")
            shutil.copyfile(txt_path, savefolder + f"/val/{txt_name}")
        else:
            shutil.copyfile(img_path, savefolder + f"/test/{img_name}")
            shutil.copyfile(txt_path, savefolder + f"/test/{txt_name}")

#振り分けの比率(train：test：val＝7：1：2)
split_data(input_dir,save_dir,0.7,0.1)

4_shuffle.pyを実行するとresourceフォルダ内のデータがランダムでcolab_uploadフォルダ内のtrain,test,valフォルダに振り分けられます。

この振り分けられたtrain,test,valフォルダをそのままgoogledriveのdataフォルダ内にuploadします。

「yolo.yaml」の作成

中身

#train,test,valのパスを指定
train: /content/drive/My Drive/data/train
val: /content/drive/My Drive/data/val
test: /content/drive/My Drive/data/test

#検出するクラスの数
nc: 1

#検出する各クラスの名前 "name1", "name2"・・・
names: ["clock"]

「yolo.yaml」をdataフォルダ内にuploadする

My Drive
    |_____data
            |_____train
            |_____val
            |_____test
            |_____yolo.yaml

colabでの操作

My Drive直下にgoogle colabのipynbファイルを作成する

ランタイムのタイプを「T4 GPU」に変更

以下をgooglecolabで実行

#ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')
%cd /content/drive/MyDrive/data

#必要なライブラリのインストール
!pip3 install ultralytics

学習開始

続いてcolabで以下を実行

!yolo task=detect mode=train model=yolov8n.pt data=yolo.yaml epochs=50 imgsz=640 batch=16

「model=」では次のうちどれかを指定します。

求める精度・速度や環境に応じて選択します。

今回は一般的なノートPCで推論するため、速度が速く軽い「YOLOv8n」を選択しています。

詳しくはultralyticsのドキュメントを見てください。

https://docs.ultralytics.com/ja/tasks/detect/

学習が終わるまで待機します

物体検出をやってみる

google driveのマイドライブある「data」フォルダ内に「runs」というフォルダが作成されています。

「/runs/detect/train/weights/best.pt」これをPCにダウンロードしてください。

コマンドプロンプトにて以下を実行

pip3 install ultralytics

Pythonで次のコードを実行（WEBカメラによるリアルタイム推論）

from ultralytics import YOLO
model = YOLO("ダウンロードしたbest.ptのパス")

#conf信頼度(conf=0.5 => 信頼度50%以上のオブジェクトを表示)
results = model(0 , show=True, conf=0.5) 
for i in enumerate(results):
    print(i)

実際のリアルタイム推論の様子（スクショ）

高い信頼度で時計として認識出来ていますね。

しかし、時計と似ているものにも反応してしまいました。

（タイマー付き電源タップ）これに関しては似すぎていると思うが...

また、取得してきた時計の画像が丸いものに偏っていたせいか、信頼度は低いものの、丸いものにも反応しています。

このような問題を回避するためには

などを行い試行錯誤する必要があります。