Databricks上のYOLOで物体検出を行う

投稿日: 投稿者:
  • 本稿では、Databricks上のNotebookで画像に写っている物体（本例ではマスク）を検出するモデルをYOLOの最新版のv8で作成して、一般的な画像からマスクを検出してみます。また、MLFlowでモデルの精度を評価します

動作環境

  • Databricks Runtime 14.3 LTS ML
  • クラスター g4dn.xlarge（16GB Memory, 1GPU）
  • YOLO
    • Ultralytics YOLOv8.1.45 🚀 Python-3.10.12 torch-2.0.1+cu118 CUDA:0 (Tesla T4, 14931MiB)
  • ソースコードはGitHubに公開しています

学習データセットの準備

  • YOLO v8でモデル学習するためのデータセットを準備します。データのフォーマットは、以下の図のように学習させたい物体ごとにクラスID, Box範囲のX軸中心、Box範囲のY軸中心、Width、Hightを1行で定義します。詳細はUltralytics YOLO formatをご確認ください。
  • 今回はYOLOフォーマットに対応したマスクを定義したデータセットをroboflowからダウンロードして利用します
  • ダウンロードしたデータをS3に配置してそのパスをDatabricksの外部ロケーションとして設定して直接Notebookから参照できるようにします
  • 以下のコードで外部ロケーションを作成します
    • Notebookパラメータとしてs3のパスを外部ロケーションパスとして設定します。設定したUnityCatalogのカタログ名、スキーマ名で外部ロケーションが作成されます
dbutils.widgets.text("external_volume_location", "")
dbutils.widgets.text("volume_catalog_name", "")
dbutils.widgets.text("volume_schema_name", "")

external_volume_location = dbutils.widgets.get("external_volume_location")
volume_catalog_name = dbutils.widgets.get("volume_catalog_name")
volume_schema_name = dbutils.widgets.get("volume_schema_name")

spark.sql(f"""
    CREATE EXTERNAL VOLUME IF NOT EXISTS {volume_catalog_name}.{volume_schema_name}.yolo8
    LOCATION "{external_volume_location}"
    COMMENT "yolo8 poc"
""")
  • 以下のようなパスで画像データを読み込むことができます
    • 以下の例ではmaskというフォルダ名がマスクのデータセット配置パスです
    • catalogとschema名は任意です
      • /Volumes/catalog/schema/yolo8/mask/train/images/xxxx.jpg
  • 学習データセットに含まれる data.yamlファイルを上記のVolumeからの絶対パスで参照できるように編集します
train: /Volumes/catalog/schema/yolo8/mask/train/images
val: /Volumes/catalog/schema/yolo8/mask/valid/images
test: /Volumes/catalog/schema/yolo8/mask/test/images

nc: 2
names: ['mask', 'no-mask']

学習（トレーニング）を実行する

  • DatabricksにYOLOv8を導入するのはとても簡単です。以下のコマンドをNotebookから実行します。他に必要なライブラリはDatabricksのMLランタイムにすべて含まれています
%pip install ultralytics

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()

dbutils.library.restartPython()
  • 学習を実行します。モデルはyolov8n.ptを利用します。YOLO v8にはあらかじめ学習済のモデルがいくつかありますが、yolov8n.ptは一番高速に動作します。詳しくは公式のパフォーマンスと精度を参照してください。
    • Notebookパラメータのdataset_locationは外部ボリュームのパスを指定します
      • 先ほど例では、/Volumes/catalog/schema/yolo8となります
dbutils.widgets.text("dataset_location", "")
dataset_location = dbutils.widgets.get("dataset_location")

if not dist.is_initialized():
    dist.init_process_group("nccl")

model = YOLO(f"{dataset_location}/yolov8n.pt")
results = model.train(
    data=f"{dataset_location}/mask/data.yaml", 
    epochs=15,
    project='/tmp/yolo8/', 
    name="yolov8_mask", 
    device=[0]
)
  • 学習回数（epochs）やその他パラメータは公式のパラメータを一覧を参考に設定してください。これらはモデルのパフォーマンスや精度に影響があります
  • nameは後述するMLFLowの実験の名前に反映されます

モデルの評価

  • YOLOv8は標準でMLFlowと統合する設定が含まれているので、特に設定せずにMLFlowにトラッキングが行われます
  • Notebookで学習の実行が終わると、右上のMLflow experimentsタブから実験結果を参照することができます。Run名が学習時に設定した名前yolov8_maskになっているので、リンクをクリックします
  • 実験結果の画面に遷移すると学習時のパラメータとメトリクスが保存されていることがわかります。上部のModel metricsタブを開くと、モデルの評価指標がグラフで可視化されて表示されます
  • 中段にmAP50-95BとmAP50Bという指標を確認することができます。mAP（mean Average Precision）とは、統計学／機械学習における（基本的に）多クラス分類のタスク（問題）に対する評価指標の一つで、各クラスのAP（Average Precision：平均適合率）を平均した値を指します。 0.0（＝0％）～1.0（＝100％）の範囲の値になり、1.0に近づくほどより良いモデルとなります

推論（物体の検出）を実行する

  • 筆者が街で撮影した画像を用いて物体（本稿の例ではマスク）を検出してみます
  • 以下のコードでNotebook上に検出対象の画像を表示させます
import cv2
from dbruntime.patches import cv2_imshow

img = cv2.imread("./detect-mask-image.png")

cv2_imshow(img)
  • この画像で推論（物体検出）を実行し、推論結果の画像を保存します
from ultralytics import YOLO
from PIL import Image

model = YOLO("/tmp/yolo8/yolov8_mask/weights/best.pt")
results = model("./detect-mask-image.png") 
for r in results:
    im_array = r.plot()  # plot a BGR numpy array of predictions
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
    im.show()  # show image
    im.save("./results.jpg")  # save image
  • 結果の画像を表示して確認してみます
  • 手前側のマスク着用と未着用が検出できているようです（奥のほうは画像があらいため、認識できていません）

まとめ

  • 本稿ではDatabricks上で画像の物体検出を実行してみました。DatabricksのMLランタイムにはあらかじめ必要となるライブラリが導入済みであることやGPUクラスタを利用しているため、環境面の設定が必要なくすぐにモデル学習のチューニングに専念することができます。YOLOは標準でMLFlowトラッキングに対応しており、モデルのチューニングに役立てることが可能です。