はじめに
YOLOv8で物体検出に続き、セグメンテーションも試してみました。
物体検出の記事はこちらになります。
環境
Python-3.8.12
torch-1.7.1+cu110
NVIDIA GeForce RTX 3090
試してみた
YOLOv8のインストール
pip install ultralytics
データセットの準備
こちらからデータを拝借
https://cocodataset.org/#download
セグメンテーションしてみる
# モデルの生成。モデルは自動でダウンロードされます。
model = YOLO("yolov8n-seg.pt")
# 推論実行
results = model("./datasets/test2014/COCO_test2014_000000515537.jpg", save=True)
モデルは以下から選ぶことができます。
また推論実行時にsave=Trueを渡すことで、結果が出力されます。
他にも、confやclassesを渡すことで、検出する物体の種類などを絞ることができます。
結果は以下になります。
検出された人物やボールの領域が赤く塗られています。
次に、resultsの中身も確認します。
results[0].cpu().boxes.numpy().boxes
こちらは、検出されたバウンディングBoxの4つの頂点座標とconfidence score、クラスを確認することができます。
results[0].masks.cpu().numpy().masks
masks変数の中には、検出された物体ごとのマスク画像データが格納されています。
試しに、画像として出力してみます。
plt.imshow(results[0].masks.cpu().numpy().masks[0])
検出された部分が1、その他は0として表示されました。
おわりに
セグメンテーションモデルを初めて触りましたが、こんなに簡単に利用できたことに驚きました。
今後はもっと色んなことに試してみたいと思います。
以上、最後まで読んで頂きありがとうございました。