DatabricksにおけるYOLOv8を用いた物体検知(実践編)

こちらを翻訳したので、ノートブックも実行してみようと思ったのですがデータがありませんでした…。

良さそうなサンプルがないかなと探したら、こちらに行き当たりました。

こちらを参考にしながらノートブックを動かしてみます。こちらのGPUクラスターを使いました。

データの準備

事前にボリュームを作成しておきます。ここでは、/Volumes/users/takaaki_yayoi/yolo/と言うボリュームにしています。

ファイルをダウンロードします。

!wget https://www.dropbox.com/s/qvglw8pqo16769f/pothole_dataset_v8.zip?dl=1 -O /Volumes/users/takaaki_yayoi/yolo/pothole_dataset_v8.zip

解凍します。

%sh
cd /Volumes/users/takaaki_yayoi/yolo
unzip pothole_dataset_v8.zip

このようになります。

セットアップ

%pip install ultralytics==8.2.53 opencv-python
%pip install ray[default]>=2.3.0
dbutils.library.restartPython()

import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

from ultralytics import YOLO

import mlflow
import torch.distributed as dist
from ultralytics import settings
from mlflow.types.schema import Schema, ColSpec
from mlflow.models.signature import ModelSignature

os.environ["OMP_NUM_THREADS"] = "4"  # 必要に応じて設定

input_schema = Schema(
    [
        ColSpec("string", "image_source"),
    ]
)
output_schema = Schema([ColSpec("string","class_name"),
                        ColSpec("integer","class_num"),
                        ColSpec("double","confidence")]
                       )

signature = ModelSignature(inputs=input_schema, 
                           outputs=output_schema)

settings.update({"mlflow":False})


##############################################################################
## モデルの結果と predict() メソッドをキャプチャするための YOLOC クラスの作成        ##
##############################################################################

class YOLOC(mlflow.pyfunc.PythonModel):
  def __init__(self, point_file):
    self.point_file=point_file

  def load_context(self, context):
    from ultralytics import YOLO

# MLflowの設定
mlflow.autolog(disable=False)
mlflow.end_run()

# プロジェクト構成のディレクトリの設定
project_location = '/Volumes/users/takaaki_yayoi/yolo/pothole_dataset_v8/'
os.makedirs(f'{project_location}/training_runs/', exist_ok=True)
os.chdir(f'{project_location}/training_runs/')

トレーニング

以下で指定するYAMLファイルを準備します。ノートブックと同じパスに保存します。

pothole_v8.yaml

path: /Volumes/users/takaaki_yayoi/yolo/pothole_dataset_v8/
train: 'train/images'
val: 'valid/images'
 
# class names
names: 
  0: 'pothole'

動作確認目的なので、エポック数は10にしています。必要に応じて増やしてください。

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


# MLflowトレーニングプロセスの起動
# このエクスペリメントにおけるバッチ、エポック数などを指定

with mlflow.start_run():
  #model = YOLO("yolov8l-seg.pt")
  model = YOLO('yolov8n.pt')

  model.train(
          batch=8,
          device=[0],
          #data=f"{project_location}/training/data.yaml",
          data="/Workspace/Users/[e-mail]/20240711_object_detection/pothole_v8.yaml",
          epochs=10,
          project='/tmp/taka_pothole/',
          exist_ok=True,
      )


  mlflow.log_params(vars(model.trainer.model.args))
  yolo_wrapper = YOLOC(model.trainer.best)
  mlflow.pyfunc.log_model(artifact_path = "model",
                          artifacts = {'model_path': str(model.trainer.save_dir), 
                                       "best_point": str(model.trainer.best)},
                          python_model = yolo_wrapper,
                          signature = signature
                          )

トレーニングが終わりました。

評価

##########
# モデルと入力画像が与えられたら境界ボックスを予測する関数
##########

def predict_on_image(model, img):
  import pandas as pd
  import json, os

  # ロードされた画像に対して境界ボックスの予測を生成するためにモデルを使用
  results = {}
  result_list = []

  model_results = model.predict(img)
  label_names, boxes, confidence = model_results[0].names, model_results[0].boxes.cpu().numpy(), model_results[0].boxes.conf.tolist()
  label_result = [int(i) for i in model_results[0].boxes.cls.tolist()]
  named_results = [{'class_name': label_names[i], 
                    'class_num': i, 
                    'confidence': confidence[count],
                    'box': boxes[count].xyxy[0].astype(int).tolist()} for count, i in enumerate(label_result)
                  ]
  results['labels'] = named_results

  return results


##########
# 指定された信頼度の閾値に基づき、穴の周辺の境界ボックスをが描かれた画像を表示する関数
##########

def prepare_img_w_box(img_in, results, confidence_thresh):

  # ボックスのポイントに変換するヘルパー関数
  def _convert_box_yolo_cv2(boxpoints):
    xmin, ymin, xmax, ymax = boxpoints
    x, y, w, h = xmin, ymin, xmax-xmin, ymax-ymin
    return x,y,w,h
  
  # それぞれの結果に対してループを行い、信頼度の条件を満たしたらボックスを配置
  for solar_panel in results['labels']:

    if solar_panel['confidence'] > confidence_thresh:
      x,y,w,h = _convert_box_yolo_cv2(solar_panel['box'])
      cv2.rectangle(img_in, (x,y), (x+w,y+h), (0,255,0), 2)

  cv2_imshow(img_in)

import cv2
from dbruntime.patches import cv2_imshow

# opencv2経由で画像をロード
img_loc = f'{project_location}/train/images/China_Drone_000205.jpg'
img = cv2.imread(img_loc)

cv2_imshow(img)

推論します。

results = predict_on_image(model, img)
results

{'labels': [{'class_name': 'pothole',
   'class_num': 0,
   'confidence': 0.3940852880477905,
   'box': [136, 73, 181, 124]},
  {'class_name': 'pothole',
   'class_num': 0,
   'confidence': 0.2777845561504364,
   'box': [326, 0, 447, 76]}]}

confidenceの値をベースに閾値を指定します。

prepare_img_w_box(img, results, 0.2)

検知できました！

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル