【Deep learning】Pythonでyolo v5を使用したトマトの分類方法 ~学習モデルの自作方法~

Labro tomatoというdatasetを使用して、トマトの分類をモデルの作成から行う方法を紹介します。下記リンクのtrainとvalのフォルダ内に学習用の画像とラベルが入っているのでこれらを使用して実際に学習させ、トマトの判別を実施するプログラムを紹介します。

https://www.kaggle.com/datasets/nexuswho/laboro-tomato

アイコン名を入力

・yoloで画像を使用したモデルの作成方法を知りたい。
・自作のモデルを作成したい。
・学習させた内容で、画像を判別する方法を知りたい。

こんな疑問に答えます。

解決できること

・yoloを使用したモデルの作成方法を理解できる。
・自分専用のモデルを作成できる。
・モデルで画像に対して、分類及びバウンディングボックス作成ができる。

現在deep learningが話題になっていますが、自分で準備した写真を使用したモデルを作成することで、製造ラインの自動欠陥検出判定を行うことができます。

今回は実際にlabro tomatoというトマトのdatasetを使用して、モデルを作成したいと思います。

🚀 0円で現役エンジニアから学べる【Techスクールオンライン】のお申込みをお勧めします。このオンラインスクールでは、現役のエンジニアから直接学ぶことができ、プログラミングの基礎から高度なスキルまでを習得できます。しかも、今なら 0円で受講できるチャンスです。エンジニア転職を考えている方やプログラミングに興味がある方、新しいスキルを習得したい方に特におすすめです。

事前準備（データのダウンロード）

最初にkaggleのサイトでトマトの画像データをダウンロードします。

https://www.kaggle.com/datasets/nexuswho/laboro-tomato

ダウンロードすると下記のように。train, valの中にimagesとlabelsが入っています。

画像のクラスは6クラスで分類され、下記クラスに分類される。big tomatoが熟し具合によって3段階、mini tomatoが熟し具合によって3段階の6クラスに分類されている。

yoloを使用するためには、事前にpytorchの環境を作成する必要があるため、設定されていない場合は、下記リンクからgit cloneでインストールしてください。

https://github.com/ultralytics/yolov5

jupyter notebookで実施するとこのような表示になります。自分の環境はpython 3.9.18を使用しています。GPUはRTX 3060Tiを使用。

import ultralytics
ultralytics.checks()

すでにGPUの環境のセッティング、pytorchのインストールが完了している場合下記表示されます。

GPUの設定がうまくいっていない場合、+cu118が表示されない場合があります。CUDAのtoolkitやpytorchのバージョン確認、pathの確認をしてください。

次にyoloのモデルをダウンロードします。

from ultralytics import YOLO
model = YOLO("yolov5s.pt")

実行するとこのようになります。

import torch

# GPUキャッシュメモリをクリア
if torch.cuda.is_available():
    torch.cuda.empty_cache()

念のため、GPUキャッシュメモリをクリアします。

次に現在のフォルダ今回私のフォルダは240131-LABROTOMATOなので、このフォルダの下の階層に、datasetsというフォルダを作成し、datasets フォルダの中にtrainフォルダ（imagesとlabels)を入れます。imagesには643枚の画像、labelsには643個の<class> <x_center> <y_center> <width> <height>が記載されています。

labelsの中身がこれです。画像に対応して、バウンディングボックスが設定されています。

x_center と y_center: バウンディングボックスの中心のx座標とy座標を、画像の幅と高さに対する比率で表します。例えば、0.5885と0.4466は、バウンディングボックスの中心が画像の幅の58.85%、高さの44.66%の位置にあることを意味します。

width と height: バウンディングボックスの幅と高さを、画像の幅と高さに対する比率で表します。これにより、バウンディングボックスが画像全体に対してどれくらいの大きさを占めるかがわかります。例えば、0.3028と0.2050は、バウンディングボックスの幅が画像の幅の30.28%、高さが画像の高さの20.50%であることを意味します。

ではフォルダの保存の階層を合わせたら、実際に学習させます。trainとvalの環境を自分のimagesの環境に合わせて、学習したモデルを作成します。

data = """
names:
  - b_fully_ripened
  - b_half_ripened
  - b_green
  - l_fully_ripened
  - l_half_ripened
  - l_green
nc: 6
path: ./
train: "C:/自分の環境/240131-labrotomato/datasets/train/images"
val: "C:/自分の環境//240131-labrotomato/datasets/val/images"
"""

# Specify the file path where you want to save the text
file_path = "dataset.yaml"

# Open the file in write mode and write the text
with open(file_path, "w") as file:
    file.write(data)

print("ready")

このプログラムを実行すると、yaml ファイルが作成されます。このファイルの意味としては、6クラスに分類し、trainフォルダ、valフォルダの位置を指定しています。

実行すると”ready”とprintされます。

その後、学習を開始します。epochs数,batchやimagszを小さくすると、時間は早くなりますがモデルの精度が落ちます。

results = model.train(data='dataset.yaml', epochs=100, imgsz=320, model="yolov5s.pt",batch=8)

実行するとこんな感じでepoch100になるまで学習していきます。私の環境で20分くらいで完了します。

完了すると、git cloneでダウンロードしたrunsのフォルダのweightsの中にbest.ptというモデルが保存されています。これが最も良いモデルとして生成されています。

また画像に対して、識別した様子がデータとして保存されます。

実際にモデルを使用して分類してみます。下記コードで、datasets/train/imagesに入っている画像に対して、バウンディングボックスとクラス名を記載してoutput2フォルダに保存を行います。

from ultralytics import YOLO
import cv2
import os
import glob

# モデルをロード
model = YOLO('runs/detect/train/weights/best.pt')

# 画像ファイルが格納されているフォルダ
image_dir = 'datasets/train/images/'
# 出力フォルダのパス
output_dir = 'output2'

# 出力フォルダが存在しない場合は作成
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 指定したフォルダから全てのjpg画像を取得
image_paths = glob.glob(os.path.join(image_dir, '*.jpg'))

for image_path in image_paths:
    # 画像を読み込む
    frame = cv2.imread(image_path)
    
    # YOLOモデルで予測
    results = model(frame)

    # 予測結果を描画
    for result in results:
        for box in result.boxes:
            # クラス名を取得
            cls_id = int(box.cls.cpu().numpy())
            cls_name = result.names[cls_id]

            # バウンディングボックスの座標を取得し、描画
            xyxy = box.xyxy.cpu().numpy()[0]
            x1, y1, x2, y2 = map(int, xyxy)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)

            # 信頼度を取得し、描画
            conf = box.conf.cpu().item()
            label = f"{cls_name} {conf:.2f}"
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 4, (255, 0, 0), 10)
    
    # 画像を出力フォルダに保存
    output_path = os.path.join(output_dir, os.path.basename(image_path))
    cv2.imwrite(output_path, frame)

print("画像の処理が完了しました。")

結果がこのようになります。こんな感じで、3つのトマトが認識されています。big-tomatoのgreenで正しく表示されています。また後ろにいるトマトもしっかり認識しています。