機械学習
記事内に商品プロモーションを含む場合があります

YOLOv8によるセグメンテーション:カスタムデータの学習から推論まで

Aru

YOLOv8ではセグメンテーション(Segmentation)タスクもサポートしています。この記事では、独自のカスタムデータセットを用いてセグメンテーションの学習と推論を行う手順について解説します。また、YOLOv8のセグメンテーションの独自の出力フォーマット(アノテーションデータ)についても詳しく解説します。

YOLOv8について

YOLOv8は、Ultralytics社が開発した物体検出モデルです。物体検出以外にもセグメンテーション、姿勢推定などにも利用できます。2023年1月に公開されたモデルということですが、モデルの進化は本当に早いと感じています。

YOLOv8のgithubのキャプチャ画像

今回はYOLOv8を用いてセグメンテーションを行ってみました。YOLOv8のセグメンテーションのデータはYOLOのフォーマットとなっていて特殊でしたので、フォーマットの確認なども行っています。

セグメンテーションタスク

セグメンテーションタスクは、与えられた画像をピクセル単位で分割し、各ピクセルが画像内のどのオブジェクトに属するかを識別するタスクです。具体的には、画像内の各ピクセルに対して、それが属する物体のクラスを推定します。

セグメンテーションでは、画像内の各ピクセルを物体クラスに関連付け、例えば「犬」「車」などを識別します。セグメンテーションは、自動運転車の環境認識、医療画像解析、農業、ロボティクス、セマンティックマッピングなどのコンピュータビジョンの多くのアプリケーションで使用されています。

セグメンテーションは、セマンティックセグメンテーションとインスタンスセグメンテーションに分かれるみたいです。

SegFormerを使ってセグメンテーションを行う記事も参考にしてください。

あわせて読みたい
SegFormerによる学習と推論:Transformerモデルでセグメンテーション
SegFormerによる学習と推論:Transformerモデルでセグメンテーション

ゼロショットセグメンテーションを実現するSAMについては以下を参考にしてください。

あわせて読みたい
Segmentation Anything(SAM)によるゼロショットセグメンテーション|使い方
Segmentation Anything(SAM)によるゼロショットセグメンテーション|使い方
あわせて読みたい
ultralytics版のSAMでゼロショットセグメンテーションに挑戦
ultralytics版のSAMでゼロショットセグメンテーションに挑戦

学習(セグメンテーション)

YOLOv8のインストール

YOLOv8のインストールは簡単です。Google Colabの場合は、以下のコマンドをセルブロックで実行すればOKです。

!pip install ultralytics

データセット(coco128-seg)

データセットはcoco128-segを利用します(https://ultralytics.com/assets/coco128-seg.zip)。このデータセットは、COCOの画像データセットのサブセットで、128枚の画像が入っています。

定義ファイルcoco128-seg.yamlには、このデータのダウンロード指定も入っており、自動的にダウンロードされますので特にダウンロードしておく必要はありません

coco128-seg.yamlは、こちらにありますが、これも学習時にダウンロードされるのでダウンロードしておく必要はありません。

YAMLファイルの中身は以下のようになっています。内容的には、train, val, testのフォルダパスと、クラス名です。downloadもオプションで指定できます。

coco128-seg.yaml
# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: yolo train data=coco128.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco128-seg  ← downloads here (7 MB)

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128-seg  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: person
  1: bicycle
  2: car
   : (略)
  77: teddy bear
  78: hair drier
  79: toothbrush

# Download script/URL (optional)
download: https://ultralytics.com/assets/coco128-seg.zip

モデルの読み込み

学習済みモデルとしては、セグメンテーションタスク向けに学習されたyolov8n-seg.ptも用意されていますが、ここでは、物体検出タスク向けに学習したyolov8n.ptをセグメンテーション用モデル(yolov8n-seg.yaml)に読み込む形とします。

学習済みモデルが学習したデータと同じCOCOデータで学習しても意味がないのでこのようにしました。

あわせて読みたい
YOLOv8(YOLOv9)で独自データを使った学習と推論を実践(物体検出)
YOLOv8(YOLOv9)で独自データを使った学習と推論を実践(物体検出)
from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n-seg.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

学習

次は、学習です。学習データの設定が書かれたファイル(coco128-seg.yaml)を指定して実行します。また、epoch数は100, 画像サイズは640を指定しました。

学習コードは以下の通りです。train()のオプションは、以下の記事または、公式サイトを参考にしてください。

results = model.train(data='coco128-seg.yaml', epochs=100, imgsz=640)

以下は、学習の経過の一部です。正しく動作すれば、このようなログが表示されます。

学習経過(一部抜粋)

学習結果を見ると、MaskのmAP50も0.701とある程度は学習できている数値になっています。

Class     Images  Instances  
all        128        929           
Box(P          R      mAP50  mAP50-95)  
 0.892      0.782      0.875      0.716
Mask(P          R      mAP50  mAP50-95)
 0.815      0.677      0.701      0.389

検証

結果の検証を行うには、以下のコードを実行します。学習時に結果は出力されているので、別途検証を行いたい場合でなければ、検証は行う必要ないかもしれません。

metrics = model.val()

戻り値のmetricsには以下のような結果が格納されています。

BOXmetrics.box.mapmAP50-95
metrics.box.map50mAP50
metrics.box.map75mAP75
metrics.box.mapsクラス毎のmAP50-95
セグメンテーションmetrics.seg.mapmAP50-95
metrics.seg.map50mAP50
metrics.seg.map75mAP75
metrics.seg.mapsクラス毎のmAP50-95

mAP(Mean Average Precision)は、セグメンテーションタスクなどのオブジェクト検出やセマンティックセグメンテーションなどのタスクにおいて、モデルの性能を評価するための一般的な指標なので、計算方法などはここでは割愛しますが、1に近いほど性能が高いことを示します。

YOLOv8のセグメンテーションのフォーマットについて

YOLO独自フォーマット形式

YOLOv8のアノテーションデータおよび、出力結果はYOLO独自のフォーマットになっています。アノテーションファイルは以下の構造になっています。

<class-index> <x1> <y1> <x2> <y2> ... <xn> <yn>

この意味は、class-indexで示されるオブジェクトが(x1, y1) → (x2, y2)→…→(xn, yn)という多角形に囲まれた領域に存在するということです。

領域のイメージ(Polygon)
YOLOv8のセグメンテーションのフォーマットのイメージ図

複数のオブジェクトがある場合は、この行が複数並びます。

出力結果も、これと同じフォーマットになります。

画像に重ねて表示させるコード

学習画像にアノテーション結果を重ねて表示させたい場合は、例えば以下のようなコードで行うことができます。

画像ファイルを読み込みます。colabでは、cv2.imshowが動作しないので、colabで実行する場合は、

from google.colab.patches import cv2_imshow

をインポートしておきます。また、cv2.im_show()の代わりにcv2_imshow()を利用します。

import cv2
from google.colab.patches import cv2_imshow

imgfile = "/content/datasets/coco128-seg/images/train2017/000000000605.jpg"
txtfile = "/content/datasets/coco128-seg/labels/train2017/000000000605.txt"

img = cv2.imread(imgfile)

次に、テキストファイルを読み込み、クラスIDをclsに、マスクをpolyに格納します。

import numpy as np
cls = []
poly = []
w, h = img.shape[1], img.shape[0]
with open(txtfile, "r") as f :
  for line in f:
    x = list(map(float, line.split()))
    cls.append(int(x[0]))
    y = np.array(x[1:]).reshape(-1,2)
    y[:,0] = y[:,0] * w
    y[:,1] = y[:,1] * h
    poly.append(y)

最後に、結果を描画します。今回は塗りつぶさずに領域を線で囲っています。

for i, e in enumerate(poly) :
  pos = e.reshape((-1,1,2)).astype(np.int32)
  cv2.polylines(img, [pos], isClosed=True, color= (255, 0, 32*i), thickness=2)
  # cv2.fillPoly(img, [pos], color= (255, 0, 32*i))

cv2_imshow(img)

結果は以下のようになります。

アノテーションデータにマスクを重畳
引用元:https://ultralytics.com/assets/coco128-seg.zip

YOLOv8の結果の可視化は、supervisonを利用すると便利です。詳細は以下の記事を参考にしてください。

YOLOの検出枠・セグメンテーション結果を「supervision」で可視化
YOLOの検出枠・セグメンテーション結果を「supervision」で可視化

予測(セグメンテーション)

学習したモデルを使って、実際に予測を行ってみます。

まず、wgetでデータを撮ってきます。wgetでUTF-8に関するエラーが出る場合は、以下のコードを実行します。

import locale
locale.getpreferredencoding = lambda: "UTF-8"

!wget https://ultralytics.com/images/bus.jpg

次に、先ほど学習したモデルを読み込み、予測を行います。

model = YOLO('/content/runs/segment/train/weights/last.pt') 
img = cv2.imread('./bus.jpg')
results = model(img, save=True, save_txt=True) 

save=True, save_txt=Trueを指定しているので、./run/predict/フォルダにラベルと結果画像が格納されているはずです。そちらも確認してください。

検出したクラス一覧を表示してみます。

# オブジェクトの種類を調べる
for e in results[0].boxes.cls.cpu():
  print(e, model.names[int(e)])
ensor(0.) person
tensor(0.) person
tensor(0.) person
tensor(0.) person
tensor(5.) bus

人が4つとバスを1つ見つけたようです。

なお、結果には、results[0].boxes, results[0].masksでアクセスできます。

results[0].boxes
ultralytics.engine.results.Boxes object with attributes:

cls: tensor([0., 0., 0., 0., 5.], device='cuda:0')
conf: tensor([0.9592, 0.9087, 0.9010, 0.4906, 0.4893], device='cuda:0')
data: tensor([[5.1545e+01, 4.0126e+02, 2.4082e+02, 8.9826e+02, 9.5918e-01, 0.0000e+00],
        [2.1993e+02, 4.0632e+02, 3.4357e+02, 8.5763e+02, 9.0866e-01, 0.0000e+00],
        [6.6932e+02, 3.8419e+02, 8.0982e+02, 8.7788e+02, 9.0104e-01, 0.0000e+00],
        [4.6940e-01, 5.5377e+02, 6.5175e+01, 8.7878e+02, 4.9061e-01, 0.0000e+00],
        [1.7728e+01, 2.2526e+02, 8.0518e+02, 7.6442e+02, 4.8932e-01, 5.0000e+00]], device='cuda:0')
          :
       (略)
results[0].masks
ultralytics.engine.results.Masks object with attributes:

data: tensor([[[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
       (略)

入力画像を表示してみます。

from google.colab.patches import cv2_imshow
import cv2
import numpy as np
cv2_imshow(img)
入力画像

検出結果を重畳してみます。

for i, e in enumerate(results[0].masks.cpu().xy) :
  pos = e.reshape((-1,1,2)).astype(np.int32)
  cv2.polylines(img, [pos], isClosed=True, color= (255, 0, 32*i), thickness=2)
  # cv2.fillPoly(img, [pos], color= (255, 0, 32*i))

cv2_imshow(img)

ある程度は正解に近づいていますが、学習が不十分にみえます。128枚の画像で100EPOCHだとこんなものかもしれません。

重畳画像
引用元:https://ultralytics.com/assets/coco128-seg.zip

おわりに

YOLOv8を使ってセグメンテーションを行ってみました。物体検出の時にも感じましたが、YOLOv8は環境がかなり整備されていて使いやすいです。その分、コード変更などがやりにくいのかもしれませんが、物体検出モデルの変更はそもそも面倒なので、そのまま利用することが多い印象です。そういう意味では、使いやすさが高いYOLOv8は、利用するモデルの選択肢として有用な気がします。

centernetを自作したことありますが、物体検出モデルはモデル自身より、枠情報→入力データや、出力データ→枠情報への変換などの、周辺処理を記述するのが面倒な印象です。

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

ABOUT ME
ある/Aru
ある/Aru
IT&機械学習エンジニア/ファイナンシャルプランナー(CFP®)
専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中
記事URLをコピーしました